import t, c
from stdint import *
import string
import stdio
import viperlib
import memhub
from linkedlist import GSList, GSListNode
from .__types import LLVMType, TypePrint, Array, PrintStructDefinition, get_struct_name, Struct, ParamNode
from .__function import Function, FunctionPrint, _ll_name_needs_quote, Param


# ============================================================
# GlobalVariable: LLVM IR 全局变量
#
# 持有名称 + 类型 + 初始化值 + 常量标志 + 链接类型。
# 继承 GSListNode[GlobalVariable] 获取强类型 Next。
# ============================================================
@t.NoVTable
class GlobalVariable(GSListNode[GlobalVariable]):
    Name: t.CChar | t.CPtr              # 全局变量名（如 ".str" / "my_global"）
    Ty: LLVMType | t.CPtr               # 变量类型
    Initializer: t.CVoid | t.CPtr       # 初始化值 Value*（None = 无初始化）
    IsConstant: t.CInt                  # 1=constant, 0=variable
    Linkage: t.CChar | t.CPtr           # 链接类型（"private"/"internal"/"external"/...）
    Section: t.CChar | t.CPtr           # 段名（None = 默认）
    UnnamedAddr: t.CInt                 # 1=unnamed_addr, 2=local_unnamed_addr, 0=none


# ============================================================
# NamedTypeNode: 命名结构体类型链表节点
#
# 持有 LLVMType* 指针（必须是命名 Struct 类型）。
# 继承 GSListNode[NamedTypeNode] 获取强类型 Next。
# ============================================================
@t.NoVTable
class NamedTypeNode(GSListNode[NamedTypeNode]):
    Ty: LLVMType | t.CPtr               # 命名结构体类型指针


# ============================================================
# LLVMModule: LLVM IR 模块容器
#
# 使用嵌入式 Head/Tail/Count 字段代替 GSList 指针，
# 避免泛型 GSList 的 __sizeof__ 返回 0 和 vtable 偏移不一致问题。
# LLVMModulePrint 输出完整 .ll 文件内容。
# 注意: 类名必须为 LLVMModule 而非 Module，避免与 ast.Module 命名冲突
# ============================================================
@t.NoVTable
class LLVMModule:
    Name: t.CChar | t.CPtr                  # 模块名
    FuncHead: Function | t.CPtr             # 函数链表头
    FuncTail: Function | t.CPtr             # 函数链表尾
    FuncCount: t.CSizeT                     # 函数计数
    GlobalHead: GlobalVariable | t.CPtr     # 全局变量链表头
    GlobalTail: GlobalVariable | t.CPtr     # 全局变量链表尾
    GlobalCount: t.CSizeT                   # 全局变量计数
    TargetTriple: t.CChar | t.CPtr          # 目标三元组（如 "x86_64-pc-linux-gnu"）
    DataLayout: t.CChar | t.CPtr            # 数据布局（如 "e-m:e-p270:32:32-...")
    NamedTypeHead: NamedTypeNode | t.CPtr   # 命名结构体类型链表头
    NamedTypeTail: NamedTypeNode | t.CPtr   # 命名结构体类型链表尾
    NamedTypeCount: t.CSizeT                # 命名结构体类型计数
    # 总计: 12 字段 * 8 = 96 字节（无 vtable）


# ============================================================
# 工厂函数
# ============================================================
def new_module(pool: memhub.MemBuddy | t.CPtr, name: t.CChar | t.CPtr) -> LLVMModule | t.CPtr:
    """创建一个模块"""
    # 硬编码 96 字节：12 字段 * 8 字节（无 vtable，@t.NoVTable）
    ptr: LLVMModule | t.CPtr = pool.alloc(96)
    if ptr is None: return None
    string.memset(ptr, 0, 96)
    ptr.Name = name
    ptr.FuncHead = None
    ptr.FuncTail = None
    ptr.FuncCount = 0
    ptr.GlobalHead = None
    ptr.GlobalTail = None
    ptr.GlobalCount = 0
    ptr.TargetTriple = None
    ptr.DataLayout = None
    ptr.NamedTypeHead = None
    ptr.NamedTypeTail = None
    ptr.NamedTypeCount = 0
    return ptr


# GlobalVariable 硬编码大小: Next(8) + Name(8) + Ty(8) + Initializer(8) + IsConstant(4) + Linkage(8) + Section(8) + UnnamedAddr(4) = 64 字节
GV_SIZE: t.CDefine = 64

# NamedTypeNode 硬编码大小: Next(8) + Ty(8) = 16 字节
NAMED_TYPE_NODE_SIZE: t.CDefine = 16

def new_global_variable(pool: memhub.MemBuddy | t.CPtr,
                        name: t.CChar | t.CPtr, ty: LLVMType | t.CPtr) -> GlobalVariable | t.CPtr:
    """创建一个全局变量"""
    ptr: GlobalVariable | t.CPtr = pool.alloc(GV_SIZE)
    if ptr is None: return None
    string.memset(ptr, 0, GV_SIZE)
    ptr.Name = name
    ptr.Ty = ty
    ptr.Initializer = None
    ptr.IsConstant = 0
    ptr.Linkage = "external"
    ptr.Section = None
    ptr.UnnamedAddr = 0
    return ptr


# ============================================================
# 链表操作（嵌入式 Head/Tail/Count，O(1) 追加）
# ============================================================
def module_add_function(mod: LLVMModule | t.CPtr, func: Function | t.CPtr):
    """将函数追加到模块函数链表（O(1)）"""
    if mod is None or func is None: return
    func.Next = None
    if mod.FuncHead is None:
        mod.FuncHead = func
    else:
        mod.FuncTail.Next = func
    mod.FuncTail = func
    mod.FuncCount += 1


def module_add_global(mod: LLVMModule | t.CPtr, gv: GlobalVariable | t.CPtr):
    """将全局变量追加到模块全局变量链表（O(1)）"""
    if mod is None or gv is None: return
    gv.Next = None
    if mod.GlobalHead is None:
        mod.GlobalHead = gv
    else:
        mod.GlobalTail.Next = gv
    mod.GlobalTail = gv
    mod.GlobalCount += 1


def module_set_target(mod: LLVMModule | t.CPtr, triple: t.CChar | t.CPtr):
    """设置目标三元组"""
    if mod is None: return
    mod.TargetTriple = triple


def module_set_datalayout(mod: LLVMModule | t.CPtr, layout: t.CChar | t.CPtr):
    """设置数据布局"""
    if mod is None: return
    mod.DataLayout = layout


# ============================================================
# 命名结构体类型管理
# ============================================================

# ============================================================
# _is_opaque_struct - 检查类型是否是 opaque 结构体
#
# opaque 结构体特征: Struct 类型，字段数为 0 且字段列表为 None。
# 通过 Struct(pool, None, 0, name) 创建，IR 输出为 %"name" = type opaque。
# 用于在 module_add_named_type 中区分 opaque 声明和完整定义，
# 确保 opaque 声明插入链表头部（在使用之前声明）。
#
# Returns: 1=是 opaque / 0=否
# ============================================================
def _is_opaque_struct(ty: LLVMType | t.CPtr) -> int:
    """检查类型是否是 opaque 结构体（无字段的命名结构体）"""
    if ty is None:
        return 0
    match ty:
        case LLVMType.Struct(fields, fcount, name):
            # opaque: 字段数为 0 且字段列表为 None
            if fcount == 0 and fields is None:
                return 1
            return 0
        case _:
            return 0


def module_add_named_type(mod: LLVMModule | t.CPtr, pool: memhub.MemBuddy | t.CPtr,
                          ty: LLVMType | t.CPtr):
    """将命名结构体类型添加到模块（用于输出 %"name" = type { ... } 定义行）

    非命名结构体（Name 为 None 或非 Struct 类型）会被忽略。
    按 Name 去重；同名时 opaque 声明可被完整定义替换（供跨模块值类型 alloca）。
    """
    if mod is None or pool is None or ty is None:
        return
    name: t.CChar | t.CPtr = get_struct_name(ty)
    if name is None:
        return
    # 去重检查（opaque 声明可被完整定义替换）
    cur: NamedTypeNode | t.CPtr = mod.NamedTypeHead
    while cur is not None:
        if cur.Ty is not None:
            existing_name: t.CChar | t.CPtr = get_struct_name(cur.Ty)
            if existing_name is not None:
                if string.strcmp(existing_name, name) == 0:
                    # 已有完整定义则无需替换
                    if _is_opaque_struct(cur.Ty) == 0:
                        return
                    # 已有 opaque，新类型也是 opaque 则不重复
                    if _is_opaque_struct(ty) != 0:
                        return
                    # 用完整定义替换 opaque 声明（原地替换 Ty 指针）
                    cur.Ty = ty
                    return
        cur = cur.Next
    # 创建新节点
    node: NamedTypeNode | t.CPtr = pool.alloc(NAMED_TYPE_NODE_SIZE)
    if node is None:
        return
    string.memset(node, 0, NAMED_TYPE_NODE_SIZE)
    node.Ty = ty
    # opaque 类型插入链表头部，确保在引用它的完整定义之前声明。
    # 例如 TypeRegistry{HashTable*, ...} 中，HashTable 的 opaque 声明
    # 必须在 TypeRegistry 定义之前，否则 llc 报 "use of undefined type"。
    # 完整定义追加到链表尾部，保持原有顺序。
    if _is_opaque_struct(ty) != 0:
        node.Next = mod.NamedTypeHead
        if mod.NamedTypeHead is None:
            mod.NamedTypeTail = node
        mod.NamedTypeHead = node
    else:
        node.Next = None
        if mod.NamedTypeHead is None:
            mod.NamedTypeHead = node
        else:
            mod.NamedTypeTail.Next = node
        mod.NamedTypeTail = node
    mod.NamedTypeCount += 1


# ============================================================
# module_ensure_opaque_for_type - 为类型中的跨模块结构体引用添加 opaque 声明
#
# 递归扫描 LLVMType 树，当遇到 Ptr(Struct(name)) 且 name 包含 '.'（SHA1 限定名）时，
# 如果该名称未在模块中定义，则添加 opaque 声明（%"name" = type opaque）。
# 这解决了跨模块结构体引用导致的 llc "use of undefined type" 错误。
# ============================================================
def _is_type_defined_in_module(mod: LLVMModule | t.CPtr, name: t.CChar | t.CPtr) -> int:
    """检查模块中是否已定义指定名称的结构体类型，返回 1=已定义 / 0=未定义"""
    if mod is None or name is None:
        return 0
    cur: NamedTypeNode | t.CPtr = mod.NamedTypeHead
    while cur is not None:
        if cur.Ty is not None:
            existing_name: t.CChar | t.CPtr = get_struct_name(cur.Ty)
            if existing_name is not None:
                if string.strcmp(existing_name, name) == 0:
                    return 1
        cur = cur.Next
    return 0


def _scan_type_for_cross_module_refs(pool: memhub.MemBuddy | t.CPtr,
                                      mod: LLVMModule | t.CPtr,
                                      ty: LLVMType | t.CPtr):
    """递归扫描类型树，为跨模块结构体引用添加 opaque 声明"""
    if ty is None or mod is None or pool is None:
        return
    match ty:
        case LLVMType.Ptr(pointee):
            if pointee is not None:
                match pointee:
                    case LLVMType.Struct(sfields, sfcount, sname):
                        if sname is not None:
                            # 检查是否为跨模块引用（名称包含 '.'）
                            if string.strchr(sname, 46) is not None:
                                # 检查是否已在模块中定义
                                if _is_type_defined_in_module(mod, sname) == 0:
                                    # 创建 opaque 类型并添加到模块
                                    opaque_ty: LLVMType | t.CPtr = Struct(pool, None, 0, sname)
                                    if opaque_ty is not None:
                                        module_add_named_type(mod, pool, opaque_ty)
                        # 递归扫描结构体字段中的跨模块引用
                        # 例如 TypeRegistry{HashTable*, MemManager*} 中的 HashTable 引用
                        if sfields is not None:
                            fld: ParamNode | t.CPtr = sfields
                            while fld is not None:
                                if fld.Ty is not None:
                                    _scan_type_for_cross_module_refs(pool, mod, fld.Ty)
                                fld = fld.Next
                    case _:
                        _scan_type_for_cross_module_refs(pool, mod, pointee)
        case LLVMType.Struct(sfields2, sfcount2, sname2):
            # 顶层结构体（非指针）：递归扫描字段
            if sfields2 is not None:
                fld2: ParamNode | t.CPtr = sfields2
                while fld2 is not None:
                    if fld2.Ty is not None:
                        _scan_type_for_cross_module_refs(pool, mod, fld2.Ty)
                    fld2 = fld2.Next
        case LLVMType.Array(elem_ty, acount):
            if elem_ty is not None:
                _scan_type_for_cross_module_refs(pool, mod, elem_ty)
        case _:
            pass


def module_ensure_opaque_for_type(mod: LLVMModule | t.CPtr,
                                   pool: memhub.MemBuddy | t.CPtr,
                                   ty: LLVMType | t.CPtr):
    """为类型中的跨模块结构体引用添加 opaque 声明

    当类型引用了 Ptr(Struct(name)) 且 name 包含 '.'（SHA1 限定名）时，
    如果该名称未在模块中定义，则添加 opaque 声明。
    """
    _scan_type_for_cross_module_refs(pool, mod, ty)


# ============================================================
# 全局变量属性设置
# ============================================================
def global_set_constant(gv: GlobalVariable | t.CPtr, is_const: t.CInt):
    """设置全局变量是否为常量"""
    if gv is None: return
    gv.IsConstant = is_const


def global_set_linkage(gv: GlobalVariable | t.CPtr, linkage: t.CChar | t.CPtr):
    """设置链接类型 (private/internal/external/... )"""
    if gv is None: return
    gv.Linkage = linkage


def global_set_unnamed_addr(gv: GlobalVariable | t.CPtr, level: t.CInt):
    """设置 unnamed_addr 属性 (0=none, 1=unnamed_addr, 2=local_unnamed_addr)"""
    if gv is None: return
    gv.UnnamedAddr = level


def global_set_section(gv: GlobalVariable | t.CPtr, section: t.CChar | t.CPtr):
    """设置段名"""
    if gv is None: return
    gv.Section = section


# ============================================================
# 输出模式常量
# ============================================================
OUTPUT_FULL: t.CDefine = 0   # 完整 IR（stub + text）
OUTPUT_STUB: t.CDefine = 1   # 仅声明（header + 类型定义 + declare + external global）
OUTPUT_TEXT: t.CDefine = 2   # 仅代码（define + global 带初值）


# ============================================================
# LLVMModulePrint: 将模块序列化为 .ll 文本
#
# mode: OUTPUT_FULL=完整, OUTPUT_STUB=仅声明, OUTPUT_TEXT=仅代码
# ============================================================
def LLVMModulePrint(buf: t.CChar | t.CPtr, size: t.CSizeT,
                mod: LLVMModule | t.CPtr, pool: memhub.MemBuddy | t.CPtr,
                mode: t.CInt):
    """将模块 mod 序列化为 .ll 文本写入 buf"""
    if mod is None or buf is None or size == 0: return

    # 头部注释（stub 和 full 模式输出，text 模式不输出）
    if mode == OUTPUT_FULL or mode == OUTPUT_STUB:
        _append_cstr(buf, size, "; ModuleID = '")
        if mod.Name is not None:
            _append_cstr(buf, size, mod.Name)
        _append_cstr(buf, size, "'\n")

        # 目标三元组
        if mod.TargetTriple is not None:
            _append_cstr(buf, size, "target triple = \"")
            _append_cstr(buf, size, mod.TargetTriple)
            _append_cstr(buf, size, "\"\n")

        # 数据布局
        if mod.DataLayout is not None:
            _append_cstr(buf, size, "target datalayout = \"")
            _append_cstr(buf, size, mod.DataLayout)
            _append_cstr(buf, size, "\"\n")

        # 空行
        _append_cstr(buf, size, "\n")

    # 命名结构体类型定义（stub 和 full 模式输出，text 模式不输出）
    if mode == OUTPUT_FULL or mode == OUTPUT_STUB:
        # 扫描所有函数签名，为跨模块结构体引用添加 opaque 定义
        # 确保 stub.ll 包含被 declare/define 引用的跨模块类型定义
        # 修复：alloc_buf 返回 viperio.Buf，但 Buf 类型未在 NamedTypeHead 中注册
        func_scan: Function | t.CPtr = mod.FuncHead
        while func_scan is not None:
            if func_scan.RetTy is not None:
                module_ensure_opaque_for_type(mod, pool, func_scan.RetTy)
            if func_scan.Params is not None:
                param_scan: Param | t.CPtr = func_scan.Params.Head
                while param_scan is not None:
                    if param_scan.Ty is not None:
                        module_ensure_opaque_for_type(mod, pool, param_scan.Ty)
                    param_scan = param_scan.Next
            func_scan = func_scan.Next
        nt: NamedTypeNode | t.CPtr = mod.NamedTypeHead
        # 结构体定义缓冲区：48 字段 × ~50 字节/字段类型名 ≈ 2400 字节，4096 足够
        # 在循环外分配一次并复用，减少 pool 碎片
        ty_buf: t.CChar | t.CPtr = pool.alloc(4096)
        while nt is not None:
            if nt.Ty is not None:
                if ty_buf is not None:
                    ty_buf[0] = '\0'
                    PrintStructDefinition(ty_buf, 4096, nt.Ty, pool)
                    _append_cstr(buf, size, ty_buf)
                    _append_cstr(buf, size, "\n")
            nt = nt.Next

        if mod.NamedTypeHead is not None:
            _append_cstr(buf, size, "\n")

    # 全局变量
    gv: GlobalVariable | t.CPtr = mod.GlobalHead
    if mode == OUTPUT_STUB:
        # stub 模式：输出 external global（无初值）
        while gv is not None:
            _print_global_stub(buf, size, gv, pool)
            _append_cstr(buf, size, "\n")
            gv = gv.Next
    else:
        # full 和 text 模式：输出完整全局变量
        # text 模式跳过无初值的全局变量（纯声明）
        gv_idx: t.CInt = 0
        while gv is not None:
            if mode == OUTPUT_TEXT and gv.Initializer is None:
                gv = gv.Next
                gv_idx += 1
                continue
            _print_global(buf, size, gv, pool)
            _append_cstr(buf, size, "\n")
            gv = gv.Next
            gv_idx += 1

    if mod.GlobalHead is not None and (mode == OUTPUT_FULL or mode == OUTPUT_STUB):
        _append_cstr(buf, size, "\n")
    elif mod.GlobalHead is not None and mode == OUTPUT_TEXT:
        # text 模式下如果输出了全局变量，也加空行
        had_global: t.CInt = 0
        gv2: GlobalVariable | t.CPtr = mod.GlobalHead
        while gv2 is not None:
            if gv2.Initializer is not None:
                had_global = 1
                break
            gv2 = gv2.Next
        if had_global != 0:
            _append_cstr(buf, size, "\n")

    # 遍历函数
    func: Function | t.CPtr = mod.FuncHead
    func_idx: t.CInt = 0
    while func is not None:
        if mode == OUTPUT_STUB:
            # stub 模式：仅输出 declare（IsDeclared==1）
            if func.IsDeclared == 1:
                FunctionPrint(buf, size, func, pool)
                _append_cstr(buf, size, "\n")
        elif mode == OUTPUT_TEXT:
            # text 模式：仅输出 define（IsDeclared==0）
            if func.IsDeclared == 0:
                FunctionPrint(buf, size, func, pool)
                _append_cstr(buf, size, "\n")
        else:
            # full 模式：输出所有函数
            FunctionPrint(buf, size, func, pool)
            _append_cstr(buf, size, "\n")
        func = func.Next
        func_idx += 1


# ============================================================
# 全局变量打印
# ============================================================
def _print_global(buf: t.CChar | t.CPtr, size: t.CSizeT,
                  gv: GlobalVariable | t.CPtr, pool: memhub.MemBuddy | t.CPtr):
    """打印单个全局变量: @<name> = [linkage] [unnamed_addr] [constant] <ty> <init>"""
    if gv is None: return

    _append_cstr(buf, size, "@")
    if gv.Name is not None:
        if _ll_name_needs_quote(gv.Name) != 0:
            _append_cstr(buf, size, "\"")
            _append_cstr(buf, size, gv.Name)
            _append_cstr(buf, size, "\"")
        else:
            _append_cstr(buf, size, gv.Name)
    _append_cstr(buf, size, " = ")

    # 链接类型（有初始值时不输出 external，因为 external + 初始值是非法 IR）
    if gv.Linkage is not None:
        if gv.Initializer is not None and string.strcmp(gv.Linkage, "external") == 0:
            pass
        else:
            _append_cstr(buf, size, gv.Linkage)
            _append_cstr(buf, size, " ")

    # unnamed_addr
    if gv.UnnamedAddr == 1:
        _append_cstr(buf, size, "unnamed_addr ")
    elif gv.UnnamedAddr == 2:
        _append_cstr(buf, size, "local_unnamed_addr ")

    # constant / global
    if gv.IsConstant == 1:
        _append_cstr(buf, size, "constant ")
    else:
        _append_cstr(buf, size, "global ")

    # 类型
    ty_buf: t.CChar | t.CPtr = pool.alloc(128)
    if ty_buf is not None:
        ty_buf[0] = '\0'
        TypePrint(ty_buf, 128, gv.Ty, pool)
        _append_cstr(buf, size, ty_buf)

    # 初始化值
    if gv.Initializer is not None:
        _append_cstr(buf, size, " ")
        _append_cstr(buf, size, gv.Initializer)

    # 段名
    if gv.Section is not None:
        _append_cstr(buf, size, ", section \"")
        _append_cstr(buf, size, gv.Section)
        _append_cstr(buf, size, "\"")


# ============================================================
# _print_global_stub - 打印全局变量的声明形式（external global，无初值）
# ============================================================
def _print_global_stub(buf: t.CChar | t.CPtr, size: t.CSizeT,
                       gv: GlobalVariable | t.CPtr, pool: memhub.MemBuddy | t.CPtr):
    """打印全局变量声明: @<name> = external global <ty>"""
    if gv is None: return

    _append_cstr(buf, size, "@")
    if gv.Name is not None:
        if _ll_name_needs_quote(gv.Name) != 0:
            _append_cstr(buf, size, "\"")
            _append_cstr(buf, size, gv.Name)
            _append_cstr(buf, size, "\"")
        else:
            _append_cstr(buf, size, gv.Name)
    _append_cstr(buf, size, " = external global ")

    # 类型
    ty_buf: t.CChar | t.CPtr = pool.alloc(128)
    if ty_buf is not None:
        ty_buf[0] = '\0'
        TypePrint(ty_buf, 128, gv.Ty, pool)
        _append_cstr(buf, size, ty_buf)


# ============================================================
# 内部辅助
# ============================================================
def _append_cstr(dst: t.CChar | t.CPtr, dst_size: t.CSizeT, src: t.CChar | t.CPtr):
    """将 C 字符串 src 追加到 dst 末尾"""
    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'