macro_rules! impl_read {
    ($t:ident, $repr:ty, $register:ty, $body:expr) => {
        impl tock_registers::interfaces::Readable for $t {
            type T = $repr;
            type R = $register;

            #[inline]
            fn get(&self) -> $repr {
                $body
            }
        }
    };
}

macro_rules! impl_write {
    ($t:ident, $repr:ty, $register:ty, $value:ident, $body:expr) => {
        impl tock_registers::interfaces::Writeable for $t {
            type T = $repr;
            type R = $register;

            #[inline]
            fn set(&self, $value: $repr) {
                $body
            }
        }
    };
}

macro_rules! cr_impl_read {
    ($t:ident, $repr:ty, $cr:ident, $register:ty) => {
        impl_read!($t, $repr, $register, {
            let value: $repr;
            unsafe {
                core::arch::asm!(
                    concat!("mov %", stringify!($cr), ", {0}"),
                    out(reg) value,
                    options(att_syntax)
                );
            }
            value
        });
    };
}

macro_rules! cr_impl_write {
    ($t:ident, $repr:ty, $cr:ident, $register:ty) => {
        impl_write!($t, $repr, $register, value, {
            unsafe {
                core::arch::asm!(
                    concat!("mov {0}, %", stringify!($cr)),
                    in(reg) value,
                    options(att_syntax)
                );
            }
        });
    };
}

macro_rules! msr_impl_read {
    ($t:ident, $addr:expr, $register:ty) => {
        impl_read!($t, u64, $register, {
            let (high, low): (u32, u32);
            unsafe {
                core::arch::asm!(
                    "rdmsr",
                    in("ecx") $addr,
                    out("eax") low,
                    out("edx") high,
                    options(att_syntax)
                );
            }
            ((high as u64) << 32) | (low as u64)
        });
    };

    ($t:ident, $addr:expr) => { msr_impl_read!($t, $addr, ()); };
}

macro_rules! msr_impl_write {
    ($t:ident, $addr:expr, $register:ty) => {
        impl_write!($t, u64, $register, value, {
            let low = value as u32;
            let high = (value >> 32) as u32;
            unsafe {
                core::arch::asm!(
                    "wrmsr",
                    in("ecx") $addr,
                    in("eax") low,
                    in("edx") high,
                    options(att_syntax)
                );
            }
        });
    };

    ($t:ident, $addr:expr) => { msr_impl_write!($t, $addr, ()); };
}

mod cr0 {
    use tock_registers::register_bitfields;

    register_bitfields! {
        usize,
        #[allow(missing_docs)]
        pub CR0 [
            PG OFFSET(31) NUMBITS(1) [],
            CD OFFSET(30) NUMBITS(1) [],
            NW OFFSET(29) NUMBITS(1) [],
            AM OFFSET(18) NUMBITS(1) [],
            WP OFFSET(16) NUMBITS(1) [],
            NE OFFSET(5) NUMBITS(1) [],
            ET OFFSET(4) NUMBITS(1) [],
            TS OFFSET(3) NUMBITS(1) [],
            EM OFFSET(2) NUMBITS(1) [],
            MP OFFSET(1) NUMBITS(1) [],
            PE OFFSET(0) NUMBITS(1) [],
        ]
    }

    pub struct Reg;

    cr_impl_read!(Reg, usize, cr0, CR0::Register);
    cr_impl_write!(Reg, usize, cr0, CR0::Register);

    /// x86-64 control register 0
    pub const CR0: Reg = Reg;
}

mod cr2 {
    use tock_registers::register_bitfields;

    register_bitfields! {
        usize,
        #[allow(missing_docs)]
        pub CR2 [
            ADDRESS OFFSET(0) NUMBITS(size_of::<usize>())
        ]
    }

    pub struct Reg;

    cr_impl_read!(Reg, usize, cr2, CR2::Register);

    /// x86 control register 2
    pub const CR2: Reg = Reg;
}

mod cr3 {
    use tock_registers::{interfaces::Writeable, register_bitfields};

    register_bitfields! {
        usize,
        #[allow(missing_docs)]
        pub CR3 [
            ADDR OFFSET(12) NUMBITS(20) [],
        ]
    }

    pub struct Reg;

    cr_impl_read!(Reg, usize, cr3, CR3::Register);
    cr_impl_write!(Reg, usize, cr3, CR3::Register);

    impl Reg {
        pub fn set_address(&self, address: usize) {
            assert_eq!(address & 0xFFF, 0);
            self.write(CR3::ADDR.val(address >> 12))
        }
    }

    /// x86-64 control register 3
    pub const CR3: Reg = Reg;
}

mod cr4 {
    use tock_registers::register_bitfields;

    register_bitfields! {
        usize,
        #[allow(missing_docs)]
        pub CR4 [
            /// If set, XSAVE and extended processor states are enabled
            OSXSAVE OFFSET(18) NUMBITS(1) [],
            /// If set, PCID is enabled
            PCIDE OFFSET(17) NUMBITS(1) [],
            /// Indicates OS support for unmasked SIMD exceptions
            OSXMMEXCPT OFFSET(10) NUMBITS(1) [],
            /// Indicates OS support for FXSAVE and FXRSTOR instructions
            OSFXSR OFFSET(9) NUMBITS(1) [],
            /// Performance-Monitoring Counter enable
            PCE OFFSET(8) NUMBITS(1) [],
            /// If set, "page global" attribute is enabled
            PGE OFFSET(7) NUMBITS(1) [],
            /// Machine Check enable
            MCE OFFSET(6) NUMBITS(1) [],
            /// Physical Address Extension (enabled if 64-bit mode)
            PAE OFFSET(5) NUMBITS(1) [],
            /// Page Size Extension (should be enabled by yboot)
            PSE OFFSET(4) NUMBITS(1) [],
            /// Debugging extensions
            DE OFFSET(3) NUMBITS(1) [],
            TSD OFFSET(2) NUMBITS(1) [],
            PVI OFFSET(1) NUMBITS(1) [],
            VME OFFSET(0) NUMBITS(1) [],
        ]
    }

    pub struct Reg;

    cr_impl_read!(Reg, usize, cr4, CR4::Register);
    cr_impl_write!(Reg, usize, cr4, CR4::Register);

    /// x86-64 control register 4
    pub const CR4: Reg = Reg;
}

mod xcr0 {
    use tock_registers::{
        interfaces::{Readable, Writeable},
        register_bitfields,
    };

    register_bitfields! {
        u64,
        #[allow(missing_docs)]
        pub XCR0 [
            /// If set, x87 FPU/MMX is enabled
            X87 OFFSET(0) NUMBITS(1) [],
            /// If set, XSAVE support for MXCSR and XMM registers is enabled
            SSE OFFSET(1) NUMBITS(1) [],
            /// If set, AVX is enabled and XSAVE supports YMM upper halves
            AVX OFFSET(2) NUMBITS(1) [],
        ]
    }

    pub struct Reg;

    impl Readable for Reg {
        type T = u64;
        type R = XCR0::Register;

        fn get(&self) -> Self::T {
            let eax: u32;
            let edx: u32;
            unsafe {
                core::arch::asm!(
                    "xgetbv",
                    in("ecx") 0,
                    out("eax") eax,
                    out("edx") edx,
                    options(att_syntax)
                );
            }
            ((edx as u64) << 32) | (eax as u64)
        }
    }

    impl Writeable for Reg {
        type T = u64;
        type R = XCR0::Register;

        fn set(&self, value: Self::T) {
            let eax = value as u32;
            let edx = (value >> 32) as u32;
            unsafe {
                core::arch::asm!(
                    "xsetbv",
                    in("ecx") 0,
                    in("eax") eax,
                    in("edx") edx,
                    options(att_syntax)
                );
            }
        }
    }

    /// Extended control register for SSE/AVX/FPU configuration
    pub const XCR0: Reg = Reg;
}

mod msr {
    pub mod ia32_kernel_gs_base {
        const ADDR: u32 = 0xC0000102;
        pub struct Reg;

        msr_impl_read!(Reg, ADDR);
        msr_impl_write!(Reg, ADDR);

        /// IA32_KERNEL_GS_BASE model-specific register. Provides the base address for %gs-relative
        /// loads/stores.
        pub const MSR_IA32_KERNEL_GS_BASE: Reg = Reg;
    }

    pub mod ia32_fs_base {
        const ADDR: u32 = 0xC0000100;
        pub struct Reg;

        msr_impl_read!(Reg, ADDR);
        msr_impl_write!(Reg, ADDR);

        /// IA32_FS_BASE model-specific register. Provides the base address for %fs-relative
        /// loads/stores.
        pub const MSR_IA32_FS_BASE: Reg = Reg;
    }

    pub mod ia32_apic_base {
        use tock_registers::{interfaces::Readable, register_bitfields};

        register_bitfields! {
            u64,
            #[allow(missing_docs)]
            #[doc = "IA32_APIC_BASE model-specific register"]
            pub MSR_IA32_APIC_BASE [
                #[doc = "Contains a virtual page number of the Local APIC base address for this processor"]
                AddressPage OFFSET(12) NUMBITS(40) [],
                #[doc = "If set, the APIC is enabled"]
                ApicEnable OFFSET(11) NUMBITS(1) [],
                #[doc = "If set, x2APIC mode is enabled"]
                ExtendedEnable OFFSET(10) NUMBITS(1) [],
                #[doc = "If set, this CPU is a bootstrap processor"]
                BootstrapCpuCore OFFSET(8) NUMBITS(1) [],
            ]
        }

        const ADDR: u32 = 0x0000001B;
        pub struct Reg;

        msr_impl_read!(Reg, ADDR, MSR_IA32_APIC_BASE::Register);
        msr_impl_write!(Reg, ADDR, MSR_IA32_APIC_BASE::Register);

        impl Reg {
            #[inline]
            pub fn read_base(&self) -> u64 {
                self.read(MSR_IA32_APIC_BASE::AddressPage) << 12
            }
        }

        /// IA32_APIC_BASE model-specific register
        pub const MSR_IA32_APIC_BASE: Reg = Reg;
    }

    pub mod ia32_sfmask {
        use tock_registers::register_bitfields;

        register_bitfields! {
            u64,
            #[allow(missing_docs)]
            pub MSR_IA32_SFMASK [
                IF OFFSET(9) NUMBITS(1) [
                    Masked = 1,
                    Unmasked = 0
                ],
                TF OFFSET(8) NUMBITS(1) [
                    Masked = 1,
                    Unmasked = 0
                ]
            ]
        }

        const ADDR: u32 = 0xC0000084;
        pub struct Reg;

        msr_impl_read!(Reg, ADDR, MSR_IA32_SFMASK::Register);
        msr_impl_write!(Reg, ADDR, MSR_IA32_SFMASK::Register);

        /// IA32_SFMASK model-specific register
        pub const MSR_IA32_SFMASK: Reg = Reg;
    }

    pub mod ia32_star {
        use tock_registers::register_bitfields;

        register_bitfields! {
            u64,
            #[allow(missing_docs)]
            pub MSR_IA32_STAR [
                SYSCALL_CS_SS OFFSET(32) NUMBITS(16) [],
                SYSRET_CS_SS OFFSET(48) NUMBITS(16) [],
            ]
        }

        const ADDR: u32 = 0xC0000081;
        pub struct Reg;

        msr_impl_read!(Reg, ADDR, MSR_IA32_STAR::Register);
        msr_impl_write!(Reg, ADDR, MSR_IA32_STAR::Register);

        /// IA32_STAR model-specific register
        pub const MSR_IA32_STAR: Reg = Reg;
    }

    pub mod ia32_lstar {
        const ADDR: u32 = 0xC0000082;
        pub struct Reg;

        msr_impl_read!(Reg, ADDR);
        msr_impl_write!(Reg, ADDR);

        /// IA32_LSTAR model-specific register
        pub const MSR_IA32_LSTAR: Reg = Reg;
    }

    pub mod ia32_efer {
        use tock_registers::register_bitfields;

        register_bitfields! {
            u64,
            #[allow(missing_docs)]
            pub MSR_IA32_EFER [
                // If set, support for SYSCALL/SYSRET instructions is enabled
                SCE OFFSET(0) NUMBITS(1) [
                    Enable = 1,
                    Disable = 0
                ]
            ]
        }

        const ADDR: u32 = 0xC0000080;
        pub struct Reg;

        msr_impl_read!(Reg, ADDR, MSR_IA32_EFER::Register);
        msr_impl_write!(Reg, ADDR, MSR_IA32_EFER::Register);

        /// IA32_EFER Extended Feature Enable model-specific Register
        pub const MSR_IA32_EFER: Reg = Reg;
    }
}

pub use cr0::CR0;
pub use cr2::CR2;
pub use cr3::CR3;
pub use cr4::CR4;
pub use msr::ia32_apic_base::MSR_IA32_APIC_BASE;
pub use msr::ia32_efer::MSR_IA32_EFER;
pub use msr::ia32_fs_base::MSR_IA32_FS_BASE;
pub use msr::ia32_kernel_gs_base::MSR_IA32_KERNEL_GS_BASE;
pub use msr::ia32_lstar::MSR_IA32_LSTAR;
pub use msr::ia32_sfmask::MSR_IA32_SFMASK;
pub use msr::ia32_star::MSR_IA32_STAR;
pub use xcr0::XCR0;

use alloc::boxed::Box;
use bytemuck::{Pod, Zeroable};
use static_assertions::const_assert_eq;

#[cfg(any(target_arch = "x86", rust_analyzer))]
#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C, align(16))]
struct FpuContextInner {
    fcw: u16,
    fsw: u16,
    ftw: u8,
    _0: u8,
    fop: u16,
    fip: u32,
    fcs: u16,
    _1: u16,

    fdp: u32,
    fds: u16,
    _2: u16,
    mxcsr: u32,
    mxcsr_mask: u32,

    mm: [u128; 8],
    xmm: [u128; 8],
    _3: [u128; 14],
}

#[cfg(any(target_arch = "x86_64", rust_analyzer))]
#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C, align(16))]
struct FpuContextInner {
    fcw: u16,
    fsw: u16,
    ftw: u8,
    _0: u8,
    fop: u16,
    fip: u64,

    fdp: u64,
    mxcsr: u32,
    mxcsr_mask: u32,

    mm: [u128; 8],
    xmm: [u128; 16],
    _1: [u128; 6],
}

const_assert_eq!(size_of::<FpuContextInner>(), 512);

#[repr(C)]
pub struct FpuContext {
    inner: Box<FpuContextInner>,
}

impl FpuContext {
    pub fn new(mask_exceptions: bool) -> Self {
        const ALL_EXCEPTIONS_MASK: u32 = (1 << 12) | (1 << 11) | (1 << 10) | (1 << 9) | (1 << 8);
        let mut inner: Box<FpuContextInner> = unsafe { Box::new_zeroed().assume_init() };
        if mask_exceptions {
            inner.mxcsr |= ALL_EXCEPTIONS_MASK;
        }
        Self { inner }
    }

    /// Stores the FPU context into the `this` pointer.
    ///
    /// # Safety
    ///
    /// It is up to the caller to ensure `this` is a valid pointer to store the FPU context in.
    pub unsafe fn store(this: *mut Self) {
        #[cfg(any(target_arch = "x86", rust_analyzer))]
        unsafe {
            core::arch::x86::_fxsave(Box::as_mut_ptr(&mut (*this).inner) as _)
        }
        #[cfg(any(target_arch = "x86_64", rust_analyzer))]
        unsafe {
            core::arch::x86_64::_fxsave64(Box::as_mut_ptr(&mut (*this).inner) as _)
        }
    }

    /// Loads the FPU with the context stored in `this` pointer.
    ///
    /// # Safety
    ///
    /// It is up to the caller to ensure `this` is a valid pointer to load the FPU context from.
    pub unsafe fn restore(this: *const Self) {
        #[cfg(any(target_arch = "x86", rust_analyzer))]
        unsafe {
            core::arch::x86::_fxrstor(Box::as_ptr(&(*this).inner) as _)
        }
        #[cfg(any(target_arch = "x86_64", rust_analyzer))]
        unsafe {
            core::arch::x86_64::_fxrstor64(Box::as_ptr(&(*this).inner) as _)
        }
    }
}
