from __future__ import annotations
import llvmlite.ir as ir


class _VolatileLoadInstr(ir.LoadInstr):
    """带 LLVM `volatile` 限定符的 load 指令。

    llvmlite 0.47 的 `ir.LoadInstr` / `ir.StoreInstr` 不支持 volatile（无该属性，
    `descr` 也不输出该关键字），故按本仓库既有做法（BaseGen.VaArgInstruction
    自定义 `descr`）派生子类，只把 `load` 关键字扩成 `load volatile`。
    输出格式与 llvmlite 原始实现逐字段一致（类型、指针、align、元数据），
    仅多一个 `volatile` 关键字；语义与 TransPyV 侧 `build_load(..., vol)` 对齐。
    """

    def descr(self, buf: list[str]) -> None:
        [val] = self.operands
        if self.align is not None:
            align = ', align %d' % (self.align)
        else:
            align = ''
        buf.append("load volatile {0}, {1} {2}{3}{4}\n".format(
            self.type,
            val.type,
            val.get_reference(),
            align,
            self._stringify_metadata(leading_comma=True),
        ))


class _VolatileStoreInstr(ir.StoreInstr):
    """带 LLVM `volatile` 限定符的 store 指令（与 _VolatileLoadInstr 同款理由）。"""

    def descr(self, buf: list[str]) -> None:
        val, ptr = self.operands
        if self.align is not None:
            align = ', align %d' % (self.align)
        else:
            align = ''
        buf.append("store volatile {0} {1}, {2} {3}{4}{5}\n".format(
            val.type,
            val.get_reference(),
            ptr.type,
            ptr.get_reference(),
            align,
            self._stringify_metadata(leading_comma=True),
        ))


class MemoryOpsMixin:

    def _register_local_heap_ptr(self, val: ir.Value, ClassName: str | None = None, VarName: str | None = None) -> None:
        if isinstance(val.type, ir.PointerType):
            pointee: ir.Type = val.type.pointee
            if isinstance(pointee, (ir.LiteralStructType, ir.IdentifiedStructType)):
                self._local_heap_ptrs.append((val, 'struct', ClassName))
            elif isinstance(pointee, ir.IntType) and pointee.width == 8:
                self._local_heap_ptrs.append((val, 'i8', ClassName))
        if VarName:
            self._var_to_heap_ptr[VarName] = val

    def _unregister_local_heap_ptr(self, val: ir.Value) -> None:
        self._local_heap_ptrs = [(v, t, cn) for v, t, cn in self._local_heap_ptrs if v is not val]

    def _emit_local_heap_frees(self) -> None:
        if not self._local_heap_ptrs or not self.builder:
            return
        deleted_objects: set[ir.Value] = set()
        for ptr, ptr_type, ClassName in self._local_heap_ptrs:
            if ptr_type == 'i8':
                Loaded_ptr: ir.Value = self.builder.load(ptr, name="load_ptr")
                null_ptr: ir.Constant = ir.Constant(ir.PointerType(ir.IntType(8)), None)
                is_null: ir.Value = self.builder.icmp_unsigned('==', Loaded_ptr, null_ptr, name="is_null")
                not_null_block: ir.Block = self.builder.append_basic_block(name="not_null")
                next_block: ir.Block = self.builder.append_basic_block(name="next")
                self.builder.cbranch(is_null, next_block, not_null_block)
                self.builder.position_at_end(not_null_block)
            if ClassName and self._has_function(f'{ClassName}.__del__'):
                cast_ptr: ir.Value
                if isinstance(ptr.type, ir.PointerType) and isinstance(ptr.type.pointee, ir.IntType) and ptr.type.pointee.width == 8:
                    cast_ptr = self.builder.bitcast(ptr, ir.PointerType(self.structs[ClassName]), name=f"cast_{ClassName}")
                else:
                    cast_ptr = ptr
                if cast_ptr not in deleted_objects:
                    self.builder.call(self._get_function(f'{ClassName}.__del__'), [cast_ptr], name=f"call_{ClassName}.__del__")
                    deleted_objects.add(cast_ptr)
            raw: ir.Value
            if ptr_type == 'struct':
                raw = self.builder.bitcast(ptr, ir.PointerType(ir.IntType(8)), name="local_free_cast")
            elif ptr_type == 'i8':
                raw = ptr
            else:
                continue
            free_type: ir.FunctionType = ir.FunctionType(ir.VoidType(), [ir.PointerType(ir.IntType(8))])
            free_func: ir.Function = self._get_or_declare_func('free', free_type)
            self.builder.call(free_func, [raw])
            if ptr_type == 'i8':
                self.builder.branch(next_block)
                self.builder.position_at_end(next_block)
        self._local_heap_ptrs = []

    def _RegisterTempPtr(self, val: ir.Value) -> None:
        if isinstance(val.type, ir.PointerType):
            pointee: ir.Type = val.type.pointee
            if isinstance(pointee, (ir.LiteralStructType, ir.IdentifiedStructType)):
                for CN, ST in self.structs.items():
                    if pointee == ST:
                        self._temp_struct_ptrs.append(val)
                        return
            elif isinstance(pointee, ir.IntType) and pointee.width == 8:
                self._temp_struct_ptrs.append(val)

    def _UnregisterTempPtr(self, val: ir.Value) -> None:
        self._temp_struct_ptrs = [t for t in self._temp_struct_ptrs if t is not val]

    def _EmitTempFrees(self) -> None:
        if not self._temp_struct_ptrs or not self.builder:
            return
        for ptr in self._temp_struct_ptrs:
            raw: ir.Value
            if isinstance(ptr.type, ir.PointerType) and isinstance(ptr.type.pointee, (ir.LiteralStructType, ir.IdentifiedStructType)):
                raw = self.builder.bitcast(ptr, ir.PointerType(ir.IntType(8)), name="free_cast")
            elif isinstance(ptr.type, ir.PointerType) and isinstance(ptr.type.pointee, ir.IntType) and ptr.type.pointee.width == 8:
                raw = ptr
            else:
                continue
            free_type: ir.FunctionType = ir.FunctionType(ir.VoidType(), [ir.PointerType(ir.IntType(8))])
            free_func: ir.Function = self._get_or_declare_func('free', free_type)
            self.builder.call(free_func, [raw])
        self._temp_struct_ptrs = []

    def _ZeroConst(self, typ: ir.Type) -> ir.Constant:
        if isinstance(typ, ir.PointerType):
            return ir.Constant(typ, None)
        if isinstance(typ, (ir.LiteralStructType, ir.IdentifiedStructType)):
            return ir.Constant(typ, None)
        if isinstance(typ, ir.ArrayType):
            return ir.Constant(typ, None)
        return ir.Constant(typ, 0)

    def _alloca(self, typ: ir.Type, name: str = '', align: int | None = None, size: ir.Value | None = None) -> ir.AllocaInstr:
        if align is None:
            align = self._get_align(typ)
        instr: ir.AllocaInstr
        if size is not None:
            instr = self.builder.alloca(typ, size=size, name=name)
        else:
            instr = self.builder.alloca(typ, name=name)
        if align:
            instr.align = align
        return instr

    def _allocaEntry(self, typ: ir.Type, name: str = '', align: int | None = None) -> ir.AllocaInstr:
        if align is None:
            align = self._get_align(typ)
        saved_block: ir.Block = self.builder.block
        entry_block: ir.Block = self.func.entry_basic_block
        instr: ir.AllocaInstr
        if entry_block.instructions:
            first_instr: ir.Instruction = entry_block.instructions[0]
            self.builder.position_before(first_instr)
            instr = self.builder.alloca(typ, name=name)
            if align:
                instr.align = align
            self.builder.position_at_end(saved_block)
        else:
            self.builder.position_at_start(entry_block)
            instr = self.builder.alloca(typ, name=name)
            if align:
                instr.align = align
            self.builder.position_at_end(saved_block)
        return instr

    def _EmitLoad(self, ptr: ir.Value, name: str = '', align: int | None = None,
                  volatile: bool = False) -> ir.Value:
        """发射一条 load；volatile=True 时走 _VolatileLoadInstr（LLVM `load volatile`）。

        类型推断完全交给 LoadInstr 自身（与 llvmlite 原始路径同一套逻辑），
        不额外兜底：推不出类型时与原生 `builder.load` 一样直接抛错。
        """
        if not volatile:
            return self.builder.load(ptr, name=name, align=align)
        instr: ir.LoadInstr = _VolatileLoadInstr(self.builder.block, ptr, name=name)
        instr.align = align
        self.builder._insert(instr)
        return instr

    def _EmitStore(self, val: ir.Value, ptr: ir.Value, align: int | None = None,
                   volatile: bool = False) -> ir.StoreInstr:
        """发射一条 store；volatile=True 时走 _VolatileStoreInstr（LLVM `store volatile`）。"""
        if not volatile:
            return self.builder.store(val, ptr, align=align)
        instr: ir.StoreInstr = _VolatileStoreInstr(self.builder.block, val, ptr)
        instr.align = align
        self.builder._insert(instr)
        return instr

    def _load(self, ptr: ir.Value, name: str = '', align: int | None = None,
              volatile: bool | None = None) -> ir.LoadInstr:
        if align is None:
            if isinstance(ptr.type, ir.PointerType):
                align = self._get_align(ptr.type.pointee)
            else:
                align = 0
        # volatile=None ⇒ 按登记名自动判定（变量自身/所指物/字段，见 _IsPtrVolatile）；
        # 显式传入 True/False 时以调用方为准。判定必须先于下方 bitcast（bitcast
        # 结果的名字不再是变量名，会丢登记名）。
        if volatile is None:
            volatile = self._IsPtrVolatile(ptr)
        if isinstance(ptr.type, ir.PointerType):
            pointee_type: ir.Type = ptr.type.pointee
            if isinstance(pointee_type, ir.IdentifiedStructType) and pointee_type.name in self.typedef_int_widths:
                width: int = self.typedef_int_widths[pointee_type.name]
                if width > 0:
                    actual_type: ir.IntType = ir.IntType(width)
                    bitcast_ptr: ir.Value = self.builder.bitcast(ptr, ir.PointerType(actual_type), name="typedef_bitcast_Load")
                    return self._EmitLoad(bitcast_ptr, name=name, align=align, volatile=volatile)
        return self._EmitLoad(ptr, name=name, align=align, volatile=volatile)

    def _Load_var(self, name: str) -> ir.Value | None:
        if name in self._direct_values:
            return self._direct_values[name]
        if name in self.variables and self.variables[name] is not None:
            VarPtr: ir.Value = self.variables[name]
            if isinstance(VarPtr, ir.GlobalVariable) and isinstance(VarPtr.type, ir.PointerType) and isinstance(VarPtr.type.pointee, (ir.IdentifiedStructType, ir.LiteralStructType, ir.ArrayType)):
                return VarPtr
            if isinstance(VarPtr.type, ir.PointerType) and isinstance(VarPtr.type.pointee, (ir.IdentifiedStructType, ir.LiteralStructType, ir.ArrayType)):
                return VarPtr
            return self._load(VarPtr, name=name)
        if name in self.global_vars and name in self.module.globals:
            GVar: ir.GlobalVariable = self.module.globals[name]
            self.variables[name] = GVar
            if isinstance(GVar, ir.Function):
                return GVar
            if isinstance(GVar.type, ir.PointerType) and isinstance(GVar.type.pointee, (ir.IdentifiedStructType, ir.LiteralStructType, ir.ArrayType)):
                return GVar
            return self._load(GVar, name=name)
        if name in self.module.globals:
            GVar = self.module.globals[name]
            self.variables[name] = GVar
            if isinstance(GVar, ir.Function):
                return GVar
            if isinstance(GVar.type, ir.PointerType) and isinstance(GVar.type.pointee, (ir.IdentifiedStructType, ir.LiteralStructType, ir.ArrayType)):
                return GVar
            return self._load(GVar, name=name)
        return None

    def _coerce_value(self, val: ir.Value, target_type: ir.Type) -> ir.Value | None:
        if val.type == target_type:
            return val
        if isinstance(target_type, ir.IdentifiedStructType) and target_type.name in self.typedef_int_widths:
            if isinstance(val.type, ir.IntType):
                return val
        if isinstance(target_type, ir.IntType) and isinstance(val.type, ir.IntType):
            if target_type.width < val.type.width:
                return self.builder.trunc(val, target_type, name="trunc_store")
            elif val.type.width == 1:
                # i1 是布尔值（来自 icmp/fcmp/__contains__ 归一化），无符号语义，
                # 必须用 zext 而非 sext，否则 i1 的 1 会被符号扩展为 -1
                return self.builder.zext(val, target_type, name="zext_bool_store")
            else:
                return self.builder.sext(val, target_type, name="sext_store")
        if isinstance(target_type, ir.PointerType) and isinstance(val.type, ir.PointerType):
            if isinstance(val, ir.Constant) and val.constant is None:
                return ir.Constant(target_type, None)
            return self.builder.bitcast(val, target_type, name="ptr_bitcast_store")
        if isinstance(target_type, ir.PointerType) and isinstance(val.type, ir.IntType):
            if val.type.width < 64:
                val = self.builder.zext(val, ir.IntType(64), name="zext_ptr_store")
            return self.builder.inttoptr(val, target_type, name="int_to_ptr_store")
        if isinstance(target_type, ir.IntType) and isinstance(val.type, ir.PointerType):
            if isinstance(val, ir.Constant) and val.constant is None:
                # null 指针常量赋给 int 字段 → 0。绝不能对 null 做 load
                # （此前 pointee.width == target.width 时直接 load null，
                #  生成 `load i32, i32* null` → ASAN access-violation）
                return ir.Constant(target_type, 0)
            # 指针 → 整数：隐式转换一律禁止（与 TransPyV coerce_to_type 同口径）。
            # 原先三条路径都是静默的：同宽直接 load 指针所指物的位模式、≤8 位 load
            # 首字节、其余 ptrtoint 截断地址——把"类型写错"藏到运行期。
            from lib.core.LLVMCG.FuncGen import FailFastError
            raise FailFastError(
                f"指针隐式转整数已禁止：不能把指针类型 {val.type} 赋给整数类型 {target_type}；"
                f"如需转换请显式 cast（如 t.CUInt64T(p) / t.CInt(p)）"
            )
        if isinstance(target_type, ir.FloatType) and isinstance(val.type, ir.DoubleType):
            return self.builder.fptrunc(val, target_type, name="fptrunc_store")
        if isinstance(target_type, ir.DoubleType) and isinstance(val.type, ir.FloatType):
            return self.builder.fpext(val, target_type, name="fpext_store")
        if isinstance(target_type, ir.IntType) and isinstance(val.type, (ir.FloatType, ir.DoubleType)):
            return self.builder.fptosi(val, target_type, name="fptosi_store")
        if isinstance(target_type, (ir.FloatType, ir.DoubleType)) and isinstance(val.type, ir.IntType):
            return self.builder.sitofp(val, target_type, name="sitofp_store")
        if isinstance(target_type, ir.ArrayType) and isinstance(val.type, ir.ArrayType):
            if target_type.count == val.type.count and target_type.element == val.type.element:
                return val
            if (isinstance(target_type.element, ir.IntType) and target_type.element.width == 8
                and isinstance(val.type.element, ir.IntType) and val.type.element.width == 8):
                if target_type.count > val.type.count:
                    padded: list[ir.Constant] = list(val.constant) if hasattr(val, 'constant') else []
                    if padded and len(padded) == val.type.count:
                        padded.extend([ir.Constant(ir.IntType(8), 0)] * (target_type.count - val.type.count))
                        return ir.Constant(target_type, padded)
                elif target_type.count < val.type.count:
                    trimmed: list[ir.Constant] = list(val.constant) if hasattr(val, 'constant') else []
                    if trimmed and len(trimmed) == val.type.count:
                        return ir.Constant(target_type, trimmed[:target_type.count])
        if isinstance(val.type, ir.PointerType) and not isinstance(target_type, ir.PointerType):
            if val.type.pointee == target_type:
                return self.builder.load(val, name="Load_struct_for_store")
            if isinstance(val.type.pointee, ir.IdentifiedStructType) and isinstance(target_type, ir.IdentifiedStructType):
                if val.type.pointee.name == target_type.name:
                    Loaded: ir.Value = self.builder.load(val, name="Load_struct_for_store")
                    if Loaded.type == target_type:
                        return Loaded
        return None

    def _store(self, val: ir.Value, ptr: ir.Value, align: int | None = None,
               volatile: bool | None = None) -> ir.StoreInstr:
        if align is None:
            if isinstance(ptr.type, ir.PointerType):
                align = self._get_align(ptr.type.pointee)
            else:
                align = 0
        # volatile=None ⇒ 按登记名自动判定（见 _IsPtrVolatile）。判定先于下方任何
        # ptr 重建（null_struct_ptr_cast 等），否则会丢登记名。
        if volatile is None:
            volatile = self._IsPtrVolatile(ptr)
        if isinstance(ptr.type, ir.PointerType):
            target_type: ir.Type = ptr.type.pointee
            if isinstance(target_type, ir.IdentifiedStructType) and target_type.name in self.typedef_int_widths:
                if isinstance(val.type, ir.IntType):
                    width: int = self.typedef_int_widths[target_type.name]
                    if width > 0:
                        actual_type: ir.IntType = ir.IntType(width)
                        bitcast_ptr: ir.Value = self.builder.bitcast(ptr, ir.PointerType(actual_type), name="typedef_bitcast")
                        if val.type.width != actual_type.width:
                            if val.type.width > actual_type.width:
                                val = self.builder.trunc(val, actual_type, name="trunc_typedef")
                            else:
                                val = self.builder.zext(val, actual_type, name="zext_typedef")
                        return self._EmitStore(val, bitcast_ptr, align=align, volatile=volatile)
            if val.type != target_type:
                if isinstance(val, ir.Constant) and val.constant is None:
                    if isinstance(target_type, ir.PointerType):
                        val = ir.Constant(target_type, None)
                    elif isinstance(target_type, ir.IdentifiedStructType):
                        val = ir.Constant(ir.PointerType(target_type), None)
                        ptr = self.builder.bitcast(ptr, ir.PointerType(ir.PointerType(target_type)), name="null_struct_ptr_cast")
                    elif isinstance(target_type, ir.IntType):
                        # null 指针常量赋给 int 字段 → 0（与 _coerce_value 一致，
                        # 防止 `store i32* null, i32*` 类型错乱 IR）
                        val = ir.Constant(target_type, 0)
                else:
                    coerced: ir.Value | None = self._coerce_value(val, target_type)
                    if coerced is not None:
                        val = coerced
                    else:
                        src_info: str = self._get_node_info()
                        raise TypeError(f"cannot store {val.type} to {ptr.type}: mismatching types{src_info}")
        return self._EmitStore(val, ptr, align=align, volatile=volatile)

    def _get_var_ptr(self, name: str) -> ir.Value | None:
        if name in self._reg_values:
            val: ir.Value = self._reg_values[name]
            var: ir.AllocaInstr = self._alloca(val.type, name=name)
            self._store(val, var)
            self.variables[name] = var
            del self._reg_values[name]
            return var
        if name in self.variables and self.variables[name] is not None:
            return self.variables[name]
        return None