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

# 导入子模块
from .__types import (
    LLVMType, ParamNode, LLVMTypeCore,
    Int1, Int8, Int16, Int32, Int64, Ptr, Func, Void, Array, Struct,
    Half, Float, Double, FP128, Label,
    new_param_node, param_list_append,
    paramnode_get_ty, paramnode_get_next, paramnode_set_ty, paramnode_set_next,
    TypePrint, PrintStructDefinition, get_struct_name,
)
from .__values import (
    Value, new_value, ConstInt, ConstNull, ConstZero, ConstFloat, SSAValue, ValuePrint,
    value_get_ty, value_get_next,
    value_set_ty, value_set_next, value_set_name, value_set_isconst,
)
from .__function import (
    Line, BasicBlock, Param, Function,
    new_line, new_basic_block, new_param, new_function,
    function_add_param, function_add_block,
    block_append_line, block_append_text,
    function_set_attrs, param_set_attrs,
    function_get_name, function_get_ret_ty, function_get_params,
    function_get_param_head, function_get_next, function_is_declared,
    param_get_name, param_get_ty, param_get_next,
    FunctionPrint,
)
from .__module import (
    LLVMModule, GlobalVariable, NamedTypeNode, new_module, new_global_variable,
    module_add_function, module_add_global, module_add_named_type,
    module_set_target, module_set_datalayout,
    global_set_constant, global_set_linkage, global_set_unnamed_addr, global_set_section,
    LLVMModulePrint, OUTPUT_FULL, OUTPUT_STUB, OUTPUT_TEXT,
    module_ensure_opaque_for_type,
)
from .__builder import (
    IRBuilder, new_builder, position_at_end,
    build_alloca, build_load, build_store,
    build_add, build_sub, build_mul, build_sdiv, build_udiv, build_srem, build_urem,
    build_and, build_or, build_xor, build_shl, build_lshr, build_ashr,
    build_fadd, build_fsub, build_fmul, build_fdiv, build_frem, build_fneg,
    build_icmp, build_fcmp,
    build_br, build_cond_br, build_ret, build_ret_void,
    build_call, build_call_indirect, build_bitcast, build_sext, build_trunc, build_zext,
    build_ptrtoint, build_inttoptr,
    build_fpext, build_fptrunc, build_fp2si, build_si2fp, build_fp2ui, build_ui2fp,
    build_gep, build_gep_array, build_gep_struct,
    build_phi, build_switch, build_select, build_unreachable,
    build_atomicrmw, build_cmpxchg, build_inline_asm,
    PhiIncoming, new_phi_incoming, SwitchCase, new_switch_case,
    ICMP_EQ, ICMP_NE, ICMP_SGT, ICMP_SGE, ICMP_SLT, ICMP_SLE,
    ICMP_UGT, ICMP_UGE, ICMP_ULT, ICMP_ULE,
    FCMP_FALSE, FCMP_OEQ, FCMP_OGT, FCMP_OGE, FCMP_OLT, FCMP_OLE,
    FCMP_ONE, FCMP_ORD, FCMP_UNO, FCMP_UEQ, FCMP_UGT, FCMP_UGE,
    FCMP_ULT, FCMP_ULE, FCMP_UNE, FCMP_TRUE,
    ATOMIC_XCHG, ATOMIC_ADD, ATOMIC_SUB, ATOMIC_AND, ATOMIC_NAND,
    ATOMIC_OR, ATOMIC_XOR, ATOMIC_MAX, ATOMIC_MIN, ATOMIC_UMAX, ATOMIC_UMIN,
    ATOMIC_ORDER_NOTATOMIC, ATOMIC_ORDER_UNORDERED, ATOMIC_ORDER_MONOTONIC,
    ATOMIC_ORDER_ACQUIRE, ATOMIC_ORDER_RELEASE, ATOMIC_ORDER_ACQ_REL,
    ATOMIC_ORDER_SEQ_CST,
)
from .__verify import (
    VerifyResult, NameEntry,
    verify_function, verify_module,
    VERIFY_OK, VERIFY_ERR_NO_TERMINATOR, VERIFY_ERR_DUPLICATE_DEF,
    VERIFY_ERR_UNDEFINED_USE, VERIFY_ERR_PHI_NOT_FIRST,
    _new_verify_result,
)


# ============================================================
# 便捷工厂: 字符串复制
# ============================================================
def _str_dup(pool: memhub.MemBuddy | t.CPtr, src: t.CChar | t.CPtr) -> t.CChar | t.CPtr:
    """从 pool 分配并复制一个 C 字符串"""
    if src is None: return None
    slen: t.CSizeT = string.strlen(src)
    buf: t.CChar | t.CPtr = pool.alloc(slen + 1)
    if buf is None: return None
    string.strcpy(buf, src)
    return buf


# ============================================================
# 便捷工厂: 创建带名称的基本块
# ============================================================
def create_block(pool: memhub.MemBuddy | t.CPtr, func: Function | t.CPtr,
                 name: t.CChar | t.CPtr) -> BasicBlock | t.CPtr:
    """创建基本块并添加到函数"""
    name_dup: t.CChar | t.CPtr = _str_dup(pool, name)
    blk: BasicBlock | t.CPtr = new_basic_block(pool, name_dup)
    if blk is not None:
        function_add_block(func, blk)
    return blk


# ============================================================
# 便捷工厂: 创建整数常量值
# ============================================================
def const_int(pool: memhub.MemBuddy | t.CPtr, bits: t.CInt,
              val: t.CInt64T) -> Value | t.CPtr:
    """创建整数常量值（自动创建类型 + 名字文本）

    参数:
        bits: 位宽 (1/8/16/32/64)
        val: 整数值
    """
    # 创建类型
    ty: LLVMType | t.CPtr = pool.alloc(LLVMType.__sizeof__())
    if ty is None: return None
    string.memset(ty, 0, LLVMType.__sizeof__())
    c.DerefAs(ty, LLVMType.Int(bits))

    # 创建名字文本
    name: t.CChar | t.CPtr = pool.alloc(32)
    if name is None: return None
    viperlib.snprintf(name, 32, "%lld", val)

    return ConstInt(pool, ty, val, name)


def const_int32(pool: memhub.MemBuddy | t.CPtr, val: t.CInt64T) -> Value | t.CPtr:
    """便捷: 创建 i32 整数常量"""
    return const_int(pool, 32, val)


def const_int64(pool: memhub.MemBuddy | t.CPtr, val: t.CInt64T) -> Value | t.CPtr:
    """便捷: 创建 i64 整数常量"""
    return const_int(pool, 64, val)


# ============================================================
# 便捷工厂: 创建函数并添加到模块
# ============================================================
def create_function(pool: memhub.MemBuddy | t.CPtr, mod: LLVMModule | t.CPtr,
                    name: t.CChar | t.CPtr, ret_ty: LLVMType | t.CPtr) -> Function | t.CPtr:
    """创建函数并添加到模块"""
    name_dup: t.CChar | t.CPtr = _str_dup(pool, name)
    func: Function | t.CPtr = new_function(pool, name_dup, ret_ty)
    if func is not None:
        module_add_function(mod, func)
    return func


def add_param(pool: memhub.MemBuddy | t.CPtr, func: Function | t.CPtr,
              ty: LLVMType | t.CPtr, name: t.CChar | t.CPtr) -> Param | t.CPtr:
    """创建参数并添加到函数"""
    name_dup: t.CChar | t.CPtr = _str_dup(pool, name)
    param: Param | t.CPtr = new_param(pool, ty, name_dup)
    if param is not None:
        function_add_param(func, param)
    return param


# ============================================================
# 便捷工厂: 创建函数声明（declare）
# ============================================================
def create_declare(pool: memhub.MemBuddy | t.CPtr, mod: LLVMModule | t.CPtr,
                   name: t.CChar | t.CPtr, ret_ty: LLVMType | t.CPtr) -> Function | t.CPtr:
    """创建函数声明（declare，无函数体）并添加到模块

    与 create_function 不同，声明不生成函数体，仅生成:
        declare <ret_ty> @<name>(<params>)
    """
    name_dup: t.CChar | t.CPtr = _str_dup(pool, name)
    func: Function | t.CPtr = new_function(pool, name_dup, ret_ty)
    if func is not None:
        func.IsDeclared = 1
        module_add_function(mod, func)
    return func


# ============================================================
# 便捷工厂: 创建字符串常量全局变量
#
# 生成 LLVM IR:
#   @.str = private unnamed_addr constant [N x i8] c"contents\00"
# ============================================================
def create_global_string(pool: memhub.MemBuddy | t.CPtr, mod: LLVMModule | t.CPtr,
                        name: t.CChar | t.CPtr, text: t.CChar | t.CPtr) -> GlobalVariable | t.CPtr:
    """创建字符串常量全局变量并添加到模块

    参数:
        name: 全局变量名（如 ".str"）
        text: 字符串内容（C 字符串，不含末尾 NUL，NUL 由本函数自动添加）

    返回: GlobalVariable 指针，其类型为 [len+1 x i8]
    """
    if text is None: return None
    slen: t.CSizeT = string.strlen(text)
    count: t.CInt = slen + 1  # 含末尾 NUL

    # 创建 i8 类型
    i8_ty: LLVMType | t.CPtr = pool.alloc(LLVMType.__sizeof__())
    if i8_ty is None: return None
    string.memset(i8_ty, 0, LLVMType.__sizeof__())
    c.DerefAs(i8_ty, LLVMType.Int(8))

    # 创建 [count x i8] 数组类型
    arr_ty: LLVMType | t.CPtr = pool.alloc(LLVMType.__sizeof__())
    if arr_ty is None: return None
    string.memset(arr_ty, 0, LLVMType.__sizeof__())
    c.DerefAs(arr_ty, LLVMType.Array(i8_ty, count))

    # 创建全局变量
    name_dup: t.CChar | t.CPtr = _str_dup(pool, name)
    gv: GlobalVariable | t.CPtr = new_global_variable(pool, name_dup, arr_ty)
    if gv is None: return None

    # 设置属性: private + unnamed_addr + constant
    global_set_linkage(gv, "private")
    global_set_unnamed_addr(gv, 1)
    global_set_constant(gv, 1)

    # 构造初始化值文本: c"contents\00"
    # LLVM IR 字符串字面量格式: c"<chars>" 其中非打印字符用 \xx 转义
    # 这里简单处理: 直接复制原始字节，末尾加 \00
    init_buf: t.CChar | t.CPtr = pool.alloc(slen + 8)  # c"..." + \00 + "
    if init_buf is None: return None
    pos: t.CSizeT = 0
    init_buf[pos] = 'c'
    pos += 1
    init_buf[pos] = '"'
    pos += 1
    i: t.CSizeT = 0
    while i < slen:
        init_buf[pos] = text[i]
        pos += 1
        i += 1
    # 末尾 NUL 字节
    init_buf[pos] = '\\'
    pos += 1
    init_buf[pos] = '0'
    pos += 1
    init_buf[pos] = '0'
    pos += 1
    init_buf[pos] = '"'
    pos += 1
    init_buf[pos] = '\0'
    gv.Initializer = init_buf

    module_add_global(mod, gv)
    return gv


# ============================================================
# 便捷工厂: 创建带初始化值的整数全局变量
# ============================================================
def create_global_int(pool: memhub.MemBuddy | t.CPtr, mod: LLVMModule | t.CPtr,
                     name: t.CChar | t.CPtr, bits: t.CInt, val: t.CInt64T) -> GlobalVariable | t.CPtr:
    """创建整数全局变量并添加到模块

    生成 LLVM IR:
        @<name> = global i<bits> <val>
    """
    # 创建整数类型
    ty: LLVMType | t.CPtr = pool.alloc(LLVMType.__sizeof__())
    if ty is None: return None
    string.memset(ty, 0, LLVMType.__sizeof__())
    c.DerefAs(ty, LLVMType.Int(bits))

    name_dup: t.CChar | t.CPtr = _str_dup(pool, name)
    gv: GlobalVariable | t.CPtr = new_global_variable(pool, name_dup, ty)
    if gv is None: return None

    # 构造初始化值文本: "42" (类型由 _print_global 单独打印)
    init_buf: t.CChar | t.CPtr = pool.alloc(48)
    if init_buf is None: return None
    viperlib.snprintf(init_buf, 48, "%lld", val)
    gv.Initializer = init_buf

    module_add_global(mod, gv)
    return gv


# ============================================================
# LLVMModuleCore: 模块/函数/值创建核心，封装 pool 避免重复传参
#
# 用法:
#   mcore: LLVMModuleCore | t.CPtr = LLVMModuleCore(pool)
#   mod: LLVMModule | t.CPtr = mcore.NewModule("test")
#   func: Function | t.CPtr = mcore.CreateFunction(mod, "add", ty_i32)
#   val: Value | t.CPtr = mcore.ConstInt32(42)
# ============================================================
@t.NoVTable
class LLVMModuleCore:
    """模块/函数/值创建核心，封装 pool 避免每次创建都传 pool"""
    Pool: memhub.MemBuddy | t.CPtr

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

    # ---- 模块操作 ----
    def NewModule(self, name: t.CChar | t.CPtr) -> LLVMModule | t.CPtr:
        return new_module(self.Pool, name)

    def CreateFunction(self, mod: LLVMModule | t.CPtr, name: t.CChar | t.CPtr,
                       ret_ty: LLVMType | t.CPtr) -> Function | t.CPtr:
        return create_function(self.Pool, mod, name, ret_ty)

    def CreateDeclare(self, mod: LLVMModule | t.CPtr, name: t.CChar | t.CPtr,
                      ret_ty: LLVMType | t.CPtr) -> Function | t.CPtr:
        """创建函数声明（declare，无函数体）"""
        return create_declare(self.Pool, mod, name, ret_ty)

    def AddParam(self, func: Function | t.CPtr, ty: LLVMType | t.CPtr,
                 name: t.CChar | t.CPtr) -> Param | t.CPtr:
        return add_param(self.Pool, func, ty, name)

    def FunctionSetAttrs(self, func: Function | t.CPtr, attrs: t.CChar | t.CPtr):
        """设置函数级属性文本"""
        function_set_attrs(func, attrs)

    def ParamSetAttrs(self, param: Param | t.CPtr, attrs: t.CChar | t.CPtr):
        """设置参数属性文本"""
        param_set_attrs(param, attrs)

    def CreateBlock(self, func: Function | t.CPtr,
                    name: t.CChar | t.CPtr) -> BasicBlock | t.CPtr:
        return create_block(self.Pool, func, name)

    def NewBuilder(self, func: Function | t.CPtr) -> IRBuilder | t.CPtr:
        return new_builder(self.Pool, func)

    # ---- 全局变量 ----
    def NewGlobal(self, mod: LLVMModule | t.CPtr, name: t.CChar | t.CPtr,
                  ty: LLVMType | t.CPtr) -> GlobalVariable | t.CPtr:
        """创建全局变量并添加到模块"""
        name_dup: t.CChar | t.CPtr = _str_dup(self.Pool, name)
        gv: GlobalVariable | t.CPtr = new_global_variable(self.Pool, name_dup, ty)
        if gv is not None:
            module_add_global(mod, gv)
        return gv

    def CreateGlobalString(self, mod: LLVMModule | t.CPtr, name: t.CChar | t.CPtr,
                           text: t.CChar | t.CPtr) -> GlobalVariable | t.CPtr:
        """创建字符串常量全局变量"""
        return create_global_string(self.Pool, mod, name, text)

    def CreateGlobalInt(self, mod: LLVMModule | t.CPtr, name: t.CChar | t.CPtr,
                       bits: t.CInt, val: t.CInt64T) -> GlobalVariable | t.CPtr:
        """创建带初始化值的整数全局变量"""
        return create_global_int(self.Pool, mod, name, bits, val)

    def GlobalSetConstant(self, gv: GlobalVariable | t.CPtr, is_const: t.CInt):
        global_set_constant(gv, is_const)

    def GlobalSetLinkage(self, gv: GlobalVariable | t.CPtr, linkage: t.CChar | t.CPtr):
        global_set_linkage(gv, linkage)

    def GlobalSetUnnamedAddr(self, gv: GlobalVariable | t.CPtr, level: t.CInt):
        global_set_unnamed_addr(gv, level)

    def GlobalSetSection(self, gv: GlobalVariable | t.CPtr, section: t.CChar | t.CPtr):
        global_set_section(gv, section)

    # ---- 值创建 ----
    def ConstInt(self, bits: t.CInt, val: t.CInt64T) -> Value | t.CPtr:
        return const_int(self.Pool, bits, val)

    def ConstInt32(self, val: t.CInt64T) -> Value | t.CPtr:
        return const_int32(self.Pool, val)

    def ConstInt64(self, val: t.CInt64T) -> Value | t.CPtr:
        return const_int64(self.Pool, val)

    def ConstFloat(self, ty: LLVMType | t.CPtr, val: t.CDouble) -> Value | t.CPtr:
        return ConstFloat(self.Pool, ty, val)

    def SSAValue(self, ty: LLVMType | t.CPtr, name: t.CChar | t.CPtr) -> Value | t.CPtr:
        return SSAValue(self.Pool, ty, name)

    def ConstNull(self, ty: LLVMType | t.CPtr, name: t.CChar | t.CPtr) -> Value | t.CPtr:
        return ConstNull(self.Pool, ty, name)

    # ---- 打印 ----
    def PrintValue(self, buf: t.CChar | t.CPtr, size: t.CSizeT, val: Value | t.CPtr):
        ValuePrint(buf, size, val, self.Pool)

    def PrintFunction(self, buf: t.CChar | t.CPtr, size: t.CSizeT, func: Function | t.CPtr):
        FunctionPrint(buf, size, func, self.Pool)

    def PrintModule(self, buf: t.CChar | t.CPtr, size: t.CSizeT, mod: LLVMModule | t.CPtr):
        LLVMModulePrint(buf, size, mod, self.Pool, OUTPUT_FULL)

    # ---- SSA 验证 ----
    def NewVerifyResult(self, buf_size: t.CSizeT) -> VerifyResult | t.CPtr:
        """创建验证结果对象（buf_size 为错误消息缓冲区大小）"""
        return _new_verify_result(self.Pool, buf_size)

    def VerifyFunction(self, func: Function | t.CPtr,
                       result: VerifyResult | t.CPtr) -> t.CInt:
        """验证单个函数的 SSA 正确性，返回错误数（0=通过）"""
        return verify_function(self.Pool, func, result)

    def VerifyModule(self, mod: LLVMModule | t.CPtr,
                     result: VerifyResult | t.CPtr) -> t.CInt:
        """验证模块中所有函数的 SSA 正确性，返回错误数（0=通过）"""
        return verify_module(self.Pool, mod, result)
