import t, c
from stdint import *
import viperlib
import string
import memhub
from linkedlist import GSListNode


# ============================================================
# ParamNode: 函数参数类型链表节点
# 继承 GSListNode[ParamNode] 获取强类型 Next: ParamNode|CPtr
# Ty 用 t.CVoid | t.CPtr 前向引用，使用时转型为 (LLVMType | t.CPtr)
#
# 注意: 必须加 @t.NoVTable（理由见 __function.py 的 Line 注释）。
# ============================================================
@t.NoVTable
class ParamNode(GSListNode[ParamNode]):
    Ty: t.CVoid | t.CPtr          # LLVMType*（前向引用，使用时转型）


# ============================================================
# LLVMType: LLVM IR 类型系统（t.REnum tagged union）
#
# 变体少、无通用字段、操作通过 match 分派 —— 契合 REnum 模型。
# 布局: { i32 __tag, <max_variant_payload> }
# ============================================================
class LLVMType(t.REnum):
    # 整数类型: i1/i8/i16/i32/i64
    class Int:
        Bits: t.CInt

    # 指针类型: pointee*
    class Ptr:
        Pointee: LLVMType | t.CPtr

    # 函数类型: ret (params)
    class Func:
        Ret: LLVMType | t.CPtr
        Params: ParamNode | t.CPtr
        PCount: t.CInt

    # 数组类型: [N x elem_ty]
    class Array:
        ElemTy: LLVMType | t.CPtr
        Count: t.CInt

    # 结构体类型: {ty1, ty2, ...} 或 %"name" = type {ty1, ty2, ...}
    # Fields 复用 ParamNode 链表（ParamNode.Ty 即字段类型）
    # Name 非 None 时为命名结构体，TypePrint 输出 %"name" 而非 {...}
    class Struct:
        Fields: ParamNode | t.CPtr
        FCount: t.CInt
        Name: t.CChar | t.CPtr

    # 浮点类型: half/float/double/fp128
    class Float:
        Bits: t.CInt               # 16=half, 32=float, 64=double, 128=fp128

    # 无值类型
    class Void:
        pass

    # 标签类型
    class Label:
        pass

    # 元数据类型
    class Metadata:
        pass


# ============================================================
# 内部分配辅助
# ============================================================
def _new_type(pool: memhub.MemBuddy | t.CPtr) -> LLVMType | t.CPtr:
    """从 mpool 分配一个零初始化的 LLVMType"""
    ptr: LLVMType | t.CPtr = pool.alloc(LLVMType.__sizeof__())
    if ptr:
        string.memset(ptr, 0, LLVMType.__sizeof__())
    return ptr


# ============================================================
# 类型工厂函数
# ============================================================
def Int1(pool: memhub.MemBuddy | t.CPtr) -> LLVMType | t.CPtr:
    ptr: LLVMType | t.CPtr = _new_type(pool)
    if ptr is None: return None
    c.DerefAs(ptr, LLVMType.Int(1))
    return ptr


def Int8(pool: memhub.MemBuddy | t.CPtr) -> LLVMType | t.CPtr:
    ptr: LLVMType | t.CPtr = _new_type(pool)
    if ptr is None: return None
    c.DerefAs(ptr, LLVMType.Int(8))
    return ptr


def Int16(pool: memhub.MemBuddy | t.CPtr) -> LLVMType | t.CPtr:
    ptr: LLVMType | t.CPtr = _new_type(pool)
    if ptr is None: return None
    c.DerefAs(ptr, LLVMType.Int(16))
    return ptr


def Int32(pool: memhub.MemBuddy | t.CPtr) -> LLVMType | t.CPtr:
    ptr: LLVMType | t.CPtr = _new_type(pool)
    if ptr is None: return None
    c.DerefAs(ptr, LLVMType.Int(32))
    return ptr


def Int64(pool: memhub.MemBuddy | t.CPtr) -> LLVMType | t.CPtr:
    ptr: LLVMType | t.CPtr = _new_type(pool)
    if ptr is None: return None
    c.DerefAs(ptr, LLVMType.Int(64))
    return ptr


def Ptr(pool: memhub.MemBuddy | t.CPtr, pointee: LLVMType | t.CPtr) -> LLVMType | t.CPtr:
    ptr: LLVMType | t.CPtr = _new_type(pool)
    if ptr is None: return None
    c.DerefAs(ptr, LLVMType.Ptr(pointee))
    return ptr


def Func(pool: memhub.MemBuddy | t.CPtr, ret: LLVMType | t.CPtr,
         params: ParamNode | t.CPtr, count: t.CInt) -> LLVMType | t.CPtr:
    ptr: LLVMType | t.CPtr = _new_type(pool)
    if ptr is None: return None
    c.DerefAs(ptr, LLVMType.Func(ret, params, count))
    return ptr


def Void(pool: memhub.MemBuddy | t.CPtr) -> LLVMType | t.CPtr:
    ptr: LLVMType | t.CPtr = _new_type(pool)
    if ptr is None: return None
    c.DerefAs(ptr, LLVMType.Void())
    return ptr


def Array(pool: memhub.MemBuddy | t.CPtr, elem_ty: LLVMType | t.CPtr,
          count: t.CInt) -> LLVMType | t.CPtr:
    """创建数组类型 [count x elem_ty]"""
    ptr: LLVMType | t.CPtr = _new_type(pool)
    if ptr is None: return None
    c.DerefAs(ptr, LLVMType.Array(elem_ty, count))
    return ptr


def Struct(pool: memhub.MemBuddy | t.CPtr, fields: ParamNode | t.CPtr,
           fcount: t.CInt, name: t.CChar | t.CPtr) -> LLVMType | t.CPtr:
    """创建结构体类型 {ty1, ty2, ...}

    fields 为 ParamNode 链表头，fcount 为字段数量。
    name 非 None 时为命名结构体（如 "sha1.ClassName"），TypePrint 输出 %"name"。
    name 为 None 时为匿名结构体，TypePrint 输出 {ty1, ty2, ...}。
    """
    ptr: LLVMType | t.CPtr = _new_type(pool)
    if ptr is None: return None
    c.DerefAs(ptr, LLVMType.Struct(fields, fcount, name))
    return ptr


def Half(pool: memhub.MemBuddy | t.CPtr) -> LLVMType | t.CPtr:
    """创建半精度浮点类型 half (16-bit)"""
    ptr: LLVMType | t.CPtr = _new_type(pool)
    if ptr is None: return None
    c.DerefAs(ptr, LLVMType.Float(16))
    return ptr


def Float(pool: memhub.MemBuddy | t.CPtr) -> LLVMType | t.CPtr:
    """创建单精度浮点类型 float (32-bit)"""
    ptr: LLVMType | t.CPtr = _new_type(pool)
    if ptr is None: return None
    c.DerefAs(ptr, LLVMType.Float(32))
    return ptr


def Double(pool: memhub.MemBuddy | t.CPtr) -> LLVMType | t.CPtr:
    """创建双精度浮点类型 double (64-bit)"""
    ptr: LLVMType | t.CPtr = _new_type(pool)
    if ptr is None: return None
    c.DerefAs(ptr, LLVMType.Float(64))
    return ptr


def FP128(pool: memhub.MemBuddy | t.CPtr) -> LLVMType | t.CPtr:
    """创建四精度浮点类型 fp128 (128-bit)"""
    ptr: LLVMType | t.CPtr = _new_type(pool)
    if ptr is None: return None
    c.DerefAs(ptr, LLVMType.Float(128))
    return ptr


def Label(pool: memhub.MemBuddy | t.CPtr) -> LLVMType | t.CPtr:
    """创建标签类型 label"""
    ptr: LLVMType | t.CPtr = _new_type(pool)
    if ptr is None: return None
    c.DerefAs(ptr, LLVMType.Label())
    return ptr


# ============================================================
# ParamNode 辅助
# ============================================================
# ParamNode 硬编码大小: Next(8) + Ty(8) = 16 字节
PNODE_SIZE: t.CDefine = 16

def new_param_node(pool: memhub.MemBuddy | t.CPtr, ty: LLVMType | t.CPtr) -> ParamNode | t.CPtr:
    """创建一个参数节点，Ty 设为给定类型指针"""
    ptr: ParamNode | t.CPtr = pool.alloc(PNODE_SIZE)
    if ptr is None: return None
    string.memset(ptr, 0, PNODE_SIZE)
    ptr.Ty = ty
    ptr.Next = None
    return ptr


def param_list_append(head: ParamNode | t.CPtr, tail: ParamNode | t.CPtr,
                      node: ParamNode | t.CPtr) -> ParamNode | t.CPtr:
    """将 node 追加到参数链表尾部，返回 head 和更新后的 tail"""
    if node is None: return head
    node.Next = None
    if head is None:
        return node
    if tail is not None:
        tail.Next = node
    return node


# ============================================================
# ParamNode 字段访问器（供其他模块绕过 stub 类型限制使用）
#
# 跨模块编译时，使用方模块只有 ParamNode 的 stub 类型，
# 导致 TransPyC 静默跳过字段访问。这些访问器在拥有完整类型定义的
# 本模块内执行，能正确访问所有字段。
# ============================================================
def paramnode_get_ty(node: ParamNode | t.CPtr) -> LLVMType | t.CPtr:
    """获取参数节点的类型"""
    if node is None:
        return None
    return (LLVMType | t.CPtr)(node.Ty)


def paramnode_get_next(node: ParamNode | t.CPtr) -> ParamNode | t.CPtr:
    """获取链表下一个参数节点"""
    if node is None:
        return None
    return node.Next


def paramnode_set_ty(node: ParamNode | t.CPtr, ty: LLVMType | t.CPtr):
    """设置参数节点的类型"""
    if node is None:
        return
    node.Ty = ty


def paramnode_set_next(node: ParamNode | t.CPtr, nxt: ParamNode | t.CPtr):
    """设置链表下一个参数节点"""
    if node is None:
        return
    node.Next = nxt


# ============================================================
# LLVMTypeCore: 类型创建核心，封装 pool 避免重复传参
#
# 用法:
#   core: LLVMTypeCore | t.CPtr = LLVMTypeCore(pool)
#   ty_i32: LLVMType | t.CPtr = core.Int32()
#   ty_ptr: LLVMType | t.CPtr = core.Ptr(ty_i32)
#   ty_void: LLVMType | t.CPtr = core.Void()
# ============================================================
@t.NoVTable
class LLVMTypeCore:
    """类型创建核心，封装 pool 避免每次创建类型都传 pool"""
    Pool: memhub.MemBuddy | t.CPtr

    def __init__(self, pool: memhub.MemBuddy | t.CPtr):
        self.Pool = pool

    def Int1(self) -> LLVMType | t.CPtr:
        return Int1(self.Pool)

    def Int8(self) -> LLVMType | t.CPtr:
        return Int8(self.Pool)

    def Int16(self) -> LLVMType | t.CPtr:
        return Int16(self.Pool)

    def Int32(self) -> LLVMType | t.CPtr:
        return Int32(self.Pool)

    def Int64(self) -> LLVMType | t.CPtr:
        return Int64(self.Pool)

    def Ptr(self, pointee: LLVMType | t.CPtr) -> LLVMType | t.CPtr:
        return Ptr(self.Pool, pointee)

    def Func(self, ret: LLVMType | t.CPtr, params: ParamNode | t.CPtr,
             count: t.CInt) -> LLVMType | t.CPtr:
        return Func(self.Pool, ret, params, count)

    def Array(self, elem_ty: LLVMType | t.CPtr, count: t.CInt) -> LLVMType | t.CPtr:
        return Array(self.Pool, elem_ty, count)

    def Struct(self, fields: ParamNode | t.CPtr, fcount: t.CInt,
               name: t.CChar | t.CPtr) -> LLVMType | t.CPtr:
        return Struct(self.Pool, fields, fcount, name)

    def Void(self) -> LLVMType | t.CPtr:
        return Void(self.Pool)

    def Half(self) -> LLVMType | t.CPtr:
        return Half(self.Pool)

    def Float(self) -> LLVMType | t.CPtr:
        return Float(self.Pool)

    def Double(self) -> LLVMType | t.CPtr:
        return Double(self.Pool)

    def FP128(self) -> LLVMType | t.CPtr:
        return FP128(self.Pool)

    def Label(self) -> LLVMType | t.CPtr:
        return Label(self.Pool)

    def NewParamNode(self, ty: LLVMType | t.CPtr) -> ParamNode | t.CPtr:
        return new_param_node(self.Pool, ty)

    # ---- 链表辅助 ----
    def ParamListAppend(self, head: ParamNode | t.CPtr, tail: ParamNode | t.CPtr,
                        node: ParamNode | t.CPtr) -> ParamNode | t.CPtr:
        """将 node 追加到参数/字段链表，返回新的 tail"""
        return param_list_append(head, tail, node)

    # ---- 打印 ----
    def PrintType(self, buf: t.CChar | t.CPtr, size: t.CSizeT, ty: LLVMType | t.CPtr):
        TypePrint(buf, size, ty, self.Pool)


# ============================================================
# TypePrint: 将 LLVMType 序列化为 IR 文本
#
# 示例:
#   Int(32)          -> "i32"
#   Int(8)           -> "i8"
#   Ptr(Int(32))     -> "i32*"
#   Void             -> "void"
#   Array(Int(32),4) -> "[4 x i32]"
#   Struct(...)      -> "{i32, i8*}"
#   Func(...)        -> "i32 (i32, i8*)"
#
# pool 用于复杂类型（Array/Struct/Func）递归打印子类型时分配临时缓冲区。
# ============================================================
def TypePrint(buf: t.CChar | t.CPtr, size: t.CSizeT,
              ty: LLVMType | t.CPtr, pool: memhub.MemBuddy | t.CPtr):
    """将类型 ty 序列化为 IR 文本写入 buf（容量 size）

    pool 用于递归打印子类型时分配临时缓冲区。
    """
    if ty is None or buf is None or size == 0:
        return
    buf[0] = '\0'
    match ty:
        case LLVMType.Int(bits):
            viperlib.snprintf(buf, size, "i%d", bits)
        case LLVMType.Ptr(pointee):
            # 先打印 pointee，再追加 '*'
            TypePrint(buf, size, pointee, pool)
            pos: t.CSizeT = string.strlen(buf)
            if pos + 1 < size:
                buf[pos] = '*'
                buf[pos + 1] = '\0'
        case LLVMType.Array(elem_ty, count):
            # [N x elem_ty]
            viperlib.snprintf(buf, size, "[%d x ", count)
            tmp: t.CChar | t.CPtr = pool.alloc(128)
            if tmp is not None:
                TypePrint(tmp, 128, elem_ty, pool)
                _strappend(buf, size, tmp)
            _strappend(buf, size, "]")
        case LLVMType.Struct(fields, fcount, name):
            if name is not None:
                # 命名结构体: %"name"
                buf[0] = '%'
                buf[1] = '"'
                buf[2] = '\0'
                _strappend(buf, size, name)
                _strappend(buf, size, "\"")
            else:
                # 匿名结构体: {ty1, ty2, ...}
                _print_struct_body(buf, size, fields, pool)
        case LLVMType.Func(ret, params, count):
            # ret (ty1, ty2, ...)
            TypePrint(buf, size, ret, pool)
            _strappend(buf, size, " (")
            cur2: ParamNode | t.CPtr = params
            first2: t.CInt = 1
            tmp3: t.CChar | t.CPtr = pool.alloc(128)
            while cur2 is not None:
                if first2 == 0:
                    _strappend(buf, size, ", ")
                if tmp3 is not None:
                    TypePrint(tmp3, 128, cur2.Ty, pool)
                    _strappend(buf, size, tmp3)
                cur2 = cur2.Next
                first2 = 0
            _strappend(buf, size, ")")
        case LLVMType.Void():
            string.strcpy(buf, "void")
        case LLVMType.Label():
            string.strcpy(buf, "label")
        case LLVMType.Metadata():
            string.strcpy(buf, "metadata")
        case LLVMType.Float(bits):
            if bits == 16:
                string.strcpy(buf, "half")
            elif bits == 32:
                string.strcpy(buf, "float")
            elif bits == 64:
                string.strcpy(buf, "double")
            elif bits == 128:
                string.strcpy(buf, "fp128")
            else:
                viperlib.snprintf(buf, size, "float%x", bits)
        case _:
            string.strcpy(buf, "void")


# ============================================================
# 内部辅助
# ============================================================
def _strappend(dst: t.CChar | t.CPtr, dst_size: t.CSizeT, src: t.CChar | t.CPtr):
    """将 src 追加到 dst 末尾（不超过 dst_size）"""
    if dst is None or src is None: return
    dlen: t.CSizeT = string.strlen(dst)
    slen: t.CSizeT = string.strlen(src)
    remain: t.CSizeT = dst_size - dlen
    if remain <= 0: return
    i: t.CSizeT = 0
    while i < slen and i + 1 < remain:
        dst[dlen + i] = src[i]
        i += 1
    dst[dlen + i] = '\0'


# ============================================================
# _print_struct_body - 打印结构体字段部分 {ty1, ty2, ...}
#
# 供 TypePrint（匿名结构体）和 PrintStructDefinition（命名结构体定义行）复用
# ============================================================
def _print_struct_body(buf: t.CChar | t.CPtr, size: t.CSizeT,
                       fields: ParamNode | t.CPtr, pool: memhub.MemBuddy | t.CPtr):
    """打印 {ty1, ty2, ...} 到 buf（追加模式，不覆盖已有内容）"""
    if buf is None or size == 0:
        return
    _strappend(buf, size, "{")
    cur: ParamNode | t.CPtr = fields
    first: t.CInt = 1
    tmp: t.CChar | t.CPtr = pool.alloc(128)
    while cur is not None:
        if first == 0:
            _strappend(buf, size, ", ")
        if tmp is not None:
            TypePrint(tmp, 128, cur.Ty, pool)
            _strappend(buf, size, tmp)
        cur = cur.Next
        first = 0
    _strappend(buf, size, "}")


# ============================================================
# PrintStructDefinition - 打印命名结构体类型定义行
#
# 输出格式: %"name" = type { ty1, ty2, ... }
# 供 LLVMModulePrint 在全局变量/函数之前输出类型定义
# ============================================================
def PrintStructDefinition(buf: t.CChar | t.CPtr, size: t.CSizeT,
                          ty: LLVMType | t.CPtr, pool: memhub.MemBuddy | t.CPtr):
    """打印命名结构体类型定义行 %"name" = type { ty1, ty2, ... }

    当 fcount==0 且 fields is None 时，输出 opaque 声明: %"name" = type opaque
    """
    if ty is None or buf is None or size == 0:
        return
    match ty:
        case LLVMType.Struct(fields, fcount, name):
            if name is None:
                return
            buf[0] = '%'
            buf[1] = '"'
            buf[2] = '\0'
            _strappend(buf, size, name)
            # opaque 类型: fcount==0 且 fields is None（跨模块前向声明）
            if fcount == 0 and fields is None:
                _strappend(buf, size, "\" = type opaque")
            else:
                _strappend(buf, size, "\" = type ")
                _print_struct_body(buf, size, fields, pool)
        case _:
            return


# ============================================================
# get_struct_name - 获取命名结构体的名称
#
# 返回 Name 字段（如 "sha1.ClassName"），非 Struct 类型或匿名 Struct 返回 None
# 供 module_add_named_type 检查是否为命名结构体及去重
# ============================================================
def get_struct_name(ty: LLVMType | t.CPtr) -> t.CChar | t.CPtr:
    """获取命名结构体的名称，非命名结构体返回 None"""
    if ty is None:
        return None
    match ty:
        case LLVMType.Struct(fields, fcount, name):
            return name
        case _:
            return None



