import t, c
from stdint import *
import memhub
import string
import llvmlite
import stdio


# ============================================================
# HandlesEnum - 枚举类型注册和成员查找
#
# 管理 t.CEnum 派生类的成员信息:
# - 枚举名 → EnumEntry（基准类型 + 成员表）
# - 成员名 → EnumMember（值、类型）
#
# 使用全局数组存储，线性查找（枚举数量通常很少）
# ============================================================

ENUM_MAX: t.CDefine = 64
ENUM_MEMBER_MAX: t.CDefine = 64


# ============================================================
# EnumMember - 枚举成员条目
# ============================================================
@t.NoVTable
class EnumMember:
    Name:  t.CChar | t.CPtr           # 成员名（字符串）
    Value: t.CInt64T                  # 成员的整数值
    Ty:    llvmlite.LLVMType | t.CPtr  # 成员的 LLVM 类型（用于混用类型场景）


# ============================================================
# EnumEntry - 枚举类型条目
# ============================================================
@t.NoVTable
class EnumEntry:
    Name:        t.CChar | t.CPtr           # 枚举类名
    BaseTy:      llvmlite.LLVMType | t.CPtr  # 基准类型（所有成员类型中最大的）
    MemberCount: int                      # 成员数量
    Members:     EnumMember | t.CPtr         # 成员数组（ENUM_MEMBER_MAX 个槽位）


# ============================================================
# 全局注册表（静态分配）
# ============================================================
_enum_table: EnumEntry | t.CPtr = None
_enum_count: int = 0


# ============================================================
# init_enum_table — 初始化枚举注册表
# ============================================================
def init_enum_table(pool: memhub.MemBuddy | t.CPtr) -> int:
    """初始化枚举注册表，返回 1 成功"""
    global _enum_table
    global _enum_count

    if _enum_table is not None:
        return 1

    entry_size: t.CSizeT = EnumEntry.__sizeof__()
    _enum_table = pool.alloc(entry_size * ENUM_MAX)
    if _enum_table is None:
        return 0
    string.memset(_enum_table, 0, entry_size * ENUM_MAX)
    _enum_count = 0
    return 1


# ============================================================
# _get_enum_entry — 获取第 i 个 EnumEntry 槽位
# ============================================================
def _get_enum_entry(i: int) -> EnumEntry | t.CPtr:
    """获取第 i 个枚举条目"""
    if _enum_table is None or i < 0 or i >= ENUM_MAX:
        return None
    entry_size: t.CSizeT = EnumEntry.__sizeof__()
    addr: t.CUInt64T = t.CUInt64T(_enum_table) + i * entry_size
    return (EnumEntry | t.CPtr)(t.CVoid(addr, t.CPtr))


# ============================================================
# _get_enum_member — 获取枚举中第 i 个 EnumMember 槽位
# ============================================================
def _get_enum_member(enum_entry: EnumEntry | t.CPtr, i: int) -> EnumMember | t.CPtr:
    """获取枚举中第 i 个成员条目"""
    if enum_entry is None or i < 0 or i >= ENUM_MEMBER_MAX:
        return None
    member_size: t.CSizeT = EnumMember.__sizeof__()
    addr: t.CUInt64T = t.CUInt64T(enum_entry.Members) + i * member_size
    return (EnumMember | t.CPtr)(t.CVoid(addr, t.CPtr))


# ============================================================
# register_enum — 注册枚举类型
#
# 返回 EnumEntry 指针，可用于添加成员
# ============================================================
def register_enum(pool: memhub.MemBuddy | t.CPtr,
                  name: str,
                  base_ty: llvmlite.LLVMType | t.CPtr) -> EnumEntry | t.CPtr:
    """注册枚举类型，返回 EnumEntry 指针"""
    if init_enum_table(pool) == 0:
        return None

    # 检查是否已注册
    existing: EnumEntry | t.CPtr = find_enum(name)
    if existing is not None:
        return existing

    if _enum_count >= ENUM_MAX:
        stdio.printf("[ENUM] table full, cannot register %s\n", name)
        return None

    entry: EnumEntry | t.CPtr = _get_enum_entry(_enum_count)
    if entry is None:
        return None

    # 分配成员数组
    member_size: t.CSizeT = EnumMember.__sizeof__()
    entry.Members = pool.alloc(member_size * ENUM_MEMBER_MAX)
    if entry.Members is None:
        return None
    string.memset(entry.Members, 0, member_size * ENUM_MEMBER_MAX)

    # 复制类名
    name_len: t.CSizeT = string.strlen(name)
    name_buf: t.CChar | t.CPtr = pool.alloc(name_len + 1)
    if name_buf is not None:
        string.strcpy(name_buf, name)
        entry.Name = name_buf

    entry.BaseTy = base_ty
    entry.MemberCount = 0

    _enum_count += 1
    return entry


# ============================================================
# add_enum_member — 向枚举添加成员
# ============================================================
def add_enum_member(pool: memhub.MemBuddy | t.CPtr,
                    enum_entry: EnumEntry | t.CPtr,
                    member_name: str,
                    member_val: t.CInt64T,
                    member_ty: llvmlite.LLVMType | t.CPtr) -> int:
    """向枚举添加成员，返回成员索引（-1 失败）"""
    if enum_entry is None or member_name is None or member_ty is None:
        return -1

    if enum_entry.MemberCount >= ENUM_MEMBER_MAX:
        stdio.printf("[ENUM] member table full for %s\n", enum_entry.Name)
        return -1

    idx: int = enum_entry.MemberCount
    me: EnumMember | t.CPtr = _get_enum_member(enum_entry, idx)
    if me is None:
        return -1

    # 复制成员名
    name_len: t.CSizeT = string.strlen(member_name)
    name_buf: t.CChar | t.CPtr = pool.alloc(name_len + 1)
    if name_buf is not None:
        string.strcpy(name_buf, member_name)
        me.Name = name_buf

    me.Value = member_val
    me.Ty = member_ty

    enum_entry.MemberCount = idx + 1
    return idx


# ============================================================
# find_enum — 按枚举类名查找
# ============================================================
def find_enum(name: str) -> EnumEntry | t.CPtr:
    """按枚举类名查找，返回 EnumEntry 或 None"""
    if name is None or _enum_table is None:
        return None
    for i in range(_enum_count):
        entry: EnumEntry | t.CPtr = _get_enum_entry(i)
        if entry is not None and entry.Name is not None:
            if string.strcmp(entry.Name, name) == 0:
                return entry
    return None


# ============================================================
# lookup_enum_member — 按枚举类名和成员名查找
# ============================================================
def lookup_enum_member(enum_name: str,
                       member_name: str) -> EnumMember | t.CPtr:
    """按枚举类名和成员名查找，返回 EnumMember 或 None"""
    if enum_name is None or member_name is None:
        return None
    entry: EnumEntry | t.CPtr = find_enum(enum_name)
    if entry is None:
        return None
    for mi in range(entry.MemberCount):
        me: EnumMember | t.CPtr = _get_enum_member(entry, mi)
        if me is not None and me.Name is not None:
            if string.strcmp(me.Name, member_name) == 0:
                return me
    return None


# ============================================================
# is_enum_class — 检查类名是否为已注册枚举
# ============================================================
def is_enum_class(name: str) -> int:
    """检查类名是否为已注册枚举，返回 1=是 / 0=否"""
    if name is None:
        return 0
    if find_enum(name) is not None:
        return 1
    return 0


# ============================================================
# get_enum_base_type — 按枚举类名获取基准类型
# ============================================================
def get_enum_base_type(enum_name: str) -> llvmlite.LLVMType | t.CPtr:
    """按枚举类名获取基准类型"""
    entry: EnumEntry | t.CPtr = find_enum(enum_name)
    if entry is not None:
        return entry.BaseTy
    return None
