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


class VaArgMixin:

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

    def _get_va_end_intrinsic(self):
        if 'llvm.va_end' in self.functions:
            return self.functions['llvm.va_end']
        ftype = ir.FunctionType(ir.VoidType(), [ir.IntType(8).as_pointer()])
        func = 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):
        if not self.builder:
            return None
        func = self._get_va_start_intrinsic()
        i8ptr_type = ir.IntType(8).as_pointer()
        if va_list_ptr.type == i8ptr_type:
            self.builder.call(func, [va_list_ptr])
        else:
            ptr = 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=16):
        if not self.builder:
            return None
        void = ir.VoidType()
        asm_type = ir.FunctionType(void, [])
        mask = -alignment
        asm_str = f"andq ${mask}, %rsp"
        asm_func = 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):
        if not self.builder:
            return None
        func = self._get_va_end_intrinsic()
        i8ptr_type = ir.IntType(8).as_pointer()
        if va_list_ptr.type == i8ptr_type:
            self.builder.call(func, [va_list_ptr])
        else:
            ptr = 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, arg_type):
        if not self.builder:
            return self._ZeroConst(arg_type)
        if not hasattr(self, '_va_arg_counter'):
            self._va_arg_counter = 0
        self._va_arg_counter += 1
        import sys
        is_x86_64 = sys.platform in ('win32', 'win64', 'windows', 'linux', 'linux2', 'darwin')
        if is_x86_64:
            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 = 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)
            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 = 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