from __future__ import annotations
import llvmlite.ir as ir
from lib.core.LLVMCG.BaseGen import VaArgInstruction


class VaArgMixin:

    def _get_va_start_intrinsic(self) -> ir.Function:
        if 'llvm.va_start' in self.functions:
            return self.functions['llvm.va_start']
        ftype: ir.FunctionType = ir.FunctionType(ir.VoidType(), [ir.IntType(8).as_pointer()])
        func: ir.Function = ir.Function(self.module, ftype, name='llvm.va_start')
        self.functions['llvm.va_start'] = func
        return func

    def _get_va_end_intrinsic(self) -> ir.Function:
        if 'llvm.va_end' in self.functions:
            return self.functions['llvm.va_end']
        ftype: ir.FunctionType = ir.FunctionType(ir.VoidType(), [ir.IntType(8).as_pointer()])
        func: ir.Function = ir.Function(self.module, ftype, name='llvm.va_end')
        self.functions['llvm.va_end'] = func
        return func

    def emit_va_start(self, va_list_ptr: ir.Value) -> None:
        if not self.builder:
            return None
        func: ir.Function = self._get_va_start_intrinsic()
        i8ptr_type: ir.PointerType = ir.IntType(8).as_pointer()
        if va_list_ptr.type == i8ptr_type:
            self.builder.call(func, [va_list_ptr])
        else:
            ptr: ir.Value = self.builder.bitcast(va_list_ptr, i8ptr_type, name="va_list_i8ptr")
            self.builder.call(func, [ptr])
        return None

    def _emit_stack_align(self, alignment: int = 16) -> None:
        if not self.builder:
            return None
        void: ir.VoidType = ir.VoidType()
        asm_type: ir.FunctionType = ir.FunctionType(void, [])
        mask: int = -alignment
        asm_str: str = f"andq $${mask}, %rsp"
        asm_func: ir.InlineAsm = ir.InlineAsm(asm_type, asm_str, "~{dirflag},~{fpsr},~{flags},~{rsp}", side_effect=True)
        self.builder.call(asm_func, [], name="align_stack")
        return None

    def emit_va_end(self, va_list_ptr: ir.Value) -> None:
        if not self.builder:
            return None
        func: ir.Function = self._get_va_end_intrinsic()
        i8ptr_type: ir.PointerType = ir.IntType(8).as_pointer()
        if va_list_ptr.type == i8ptr_type:
            self.builder.call(func, [va_list_ptr])
        else:
            ptr: ir.Value = self.builder.bitcast(va_list_ptr, i8ptr_type, name="va_list_i8ptr")
            self.builder.call(func, [ptr])
        return None

    def emit_va_arg(self, va_list_ptr: ir.Value, arg_type: ir.Type) -> ir.Value:
        if not self.builder:
            return self._ZeroConst(arg_type)
        if not hasattr(self, '_va_arg_counter'):
            self._va_arg_counter: int = 0
        self._va_arg_counter += 1
        is_x86_64: bool = self._platform_info['is_x86_64']
        if is_x86_64:
            va_arg_type: ir.Type
            if isinstance(arg_type, ir.PointerType):
                va_arg_type = ir.IntType(64)
            elif isinstance(arg_type, (ir.FloatType, ir.DoubleType)):
                va_arg_type = ir.DoubleType()
            elif isinstance(arg_type, ir.IntType) and arg_type.width < 64:
                va_arg_type = ir.IntType(64)
            else:
                va_arg_type = arg_type
            instr: ir.Instruction = VaArgInstruction(self.builder.block, va_arg_type, 'va_arg', (va_list_ptr,), name=f'va_arg_raw_{self._va_arg_counter}')
            self.builder._insert(instr)
            result: ir.Value
            if va_arg_type != arg_type:
                if isinstance(arg_type, ir.PointerType):
                    result = self.builder.inttoptr(instr, arg_type, name=f'va_arg_ptr_{self._va_arg_counter}')
                elif isinstance(arg_type, ir.FloatType):
                    result = self.builder.fptrunc(instr, arg_type, name=f'va_arg_fptrunc_{self._va_arg_counter}')
                else:
                    result = self.builder.trunc(instr, arg_type, name=f'va_arg_trunc_{self._va_arg_counter}')
            else:
                result = instr
            return result
        else:
            if isinstance(arg_type, ir.IntType) and arg_type.width < 64:
                arg_type = ir.IntType(64)
            instr: ir.Instruction = VaArgInstruction(self.builder.block, arg_type, 'va_arg', (va_list_ptr,), name=f'va_arg_raw_{self._va_arg_counter}')
            self.builder._insert(instr)
            return instr