import cpu


# FPU上下文结构体
@__attribute__(aligned(16))
class fpu_context:
    data: [512]uint8 # FXSAVE区域大小


class _FPUObject:
    CPUObject: cpu._CPUObject*
    # 初始化FPU
    def __init__(self, CPUObject: cpu._CPUObject*):
        self.enable()
        self.CPUObject = CPUObject
        __asm__("fninit", [ASM_DESCR.CLOBBER_MEMORY])
    # 启用FPU
    def enable(self):
        cr0: uint64 = self.CPUObject.readCR0()
        # 清除EM位（禁用仿真），设置MP位（监控协处理器）
        cr0 &= ~(1 << 2); # 清除EM位
        cr0 |= (1 << 1);  # 设置MP位
        # 写入CR0寄存器
        self.CPUObject.writeCR0(cr0)
        # 读取CR4寄存器
        cr4: uint64 = self.CPUObject.readCR4()
        # 设置OSFXSR位和OSXMMEXCPT位，启用SSE
        cr4 |= (1 << 9) | (1 << 10)
        # 写入CR4寄存器
        self.CPUObject.writeCR4(cr4)
        # 初始化FPU
        __asm__("fninit", [ASM_DESCR.CLOBBER_MEMORY])
    # 禁用FPU
    def disable(self):
        cr0: uint64 = self.CPUObject.readCR0()
        # 设置EM位（启用仿真），清除MP位（禁用监控协处理器）
        cr0 |= (1 << 2)  # 设置EM位
        cr0 &= ~(1 << 1) # 清除MP位
        # 写入CR0寄存器
        self.CPUObject.writeCR0(cr0)
    # 保存FPU上下文
    def save_context(self, context: fpu_context*):
        # 保存FPU状态到内存
        __asm__(f"fxsave [{asminp(context, ASM_DESCR.REG_ANY)}]", [ASM_DESCR.CLOBBER_MEMORY])
    # 恢复FPU上下文
    def restore_context(self, context: const fpu_context*):
        # 从内存恢复FPU状态
        __asm__(f"fxrstor [{asminp(context, ASM_DESCR.REG_ANY)}]", [ASM_DESCR.CLOBBER_MEMORY])