#![allow(clippy::identity_op)]

use core::fmt::Debug;

use bit_field::BitField;
use bitflags::bitflags;

pub struct CpuidResult {
    pub eax: u32,
    pub ebx: u32,
    pub ecx: u32,
    pub edx: u32,
}

pub fn cpuid(eax: u32, ecx: u32) -> CpuidResult {
    let result = core::arch::x86_64::__cpuid_count(eax, ecx);

    CpuidResult {
        eax: result.eax,
        ebx: result.ebx,
        ecx: result.ecx,
        edx: result.edx,
    }
}

pub struct CpuId {
    pub max_eax: u32,
    pub vendor_id: [u8; 12],
}

pub struct Leaf1(CpuidResult);

impl Leaf1 {
    pub fn get() -> Self {
        let result = cpuid(1, 0);
        Self(result)
    }
    pub fn brand_index(&self) -> u8 {
        self.0.ebx.get_bits(0..8) as u8
    }
    pub fn clflush_line_size(&self) -> u8 {
        self.0.ebx.get_bits(8..16) as u8
    }
    pub fn max_logical_processors(&self) -> u8 {
        self.0.ebx.get_bits(16..24) as u8
    }
    pub fn initial_apic_id(&self) -> u8 {
        self.0.ebx.get_bits(24..32) as u8
    }

    pub fn family_id(&self) -> u8 {
        self.0.eax.get_bits(8..12) as u8 | ((self.0.eax.get_bits(20..28) as u8) << 4)
    }
    pub fn model_id(&self) -> u8 {
        self.0.eax.get_bits(4..8) as u8 | ((self.0.eax.get_bits(16..20) as u8) << 4)
    }
    pub fn stepping_id(&self) -> u8 {
        self.0.eax.get_bits(0..4) as u8
    }
    pub fn processor_type(&self) -> u8 {
        self.0.eax.get_bits(12..14) as u8
    }
    pub fn flags(&self) -> Leaf1Flags {
        Leaf1Flags::from_bits_truncate(self.0.ecx as u64 | ((self.0.edx as u64) << 32))
    }
}

bitflags! {
    pub struct Leaf1Flags: u64 {
        const SSE3 = 1 << 0;
        const PCLMULQDQ = 1 << 1;
        const DTES64 = 1 << 2;
        const MONITOR = 1 << 3;
        const DS_CPL = 1 << 4;
        const VMX = 1 << 5;
        const SMX = 1 << 6;
        const EIST = 1 << 7;
        const TM2 = 1 << 8;
        const SSSE3 = 1 << 9;
        const CNXT_ID = 1 << 10;
        const SDBG = 1 << 11;
        const FMA = 1 << 12;
        const CMPXCHG16B = 1 << 13;
        const XTPR_UPDATE_CONTROL = 1 << 14;
        const PDCM = 1 << 15;
        const PCID = 1 << 17;
        const DCA = 1 << 18;
        const SSE4_1 = 1 << 19;
        const SSE4_2 = 1 << 20;
        const X2APIC = 1 << 21;
        const MOVBE = 1 << 22;
        const POPCNT = 1 << 23;
        const TSC_DEADLINE_TIMER = 1 << 24;
        const AESNI = 1 << 25;
        const XSAVE = 1 << 26;
        const OSXSAVE = 1 << 27;
        const AVX = 1 <<28 ;
        const F16C = 1 << 29;
        const RDRAND = 1 << 30;

        const FPU = 1 << (0 + u32::BITS);
        const VME = 1 << (1 + u32::BITS);
        const DE  = 1 << (2 + u32::BITS);
        const PSE = 1 << (3 + u32::BITS);
        const TSC = 1 << (4 + u32::BITS);
        const MSR = 1 << (5 + u32::BITS);
        const PAE = 1 << (6 + u32::BITS);
        const MCE = 1 << (7 + u32::BITS);
        const CX8 = 1 << (8 + u32::BITS);
        const APIC = 1 << (9 + u32::BITS);
        const SEP = 1 << (11 + u32::BITS);
        const MTRR = 1 << (12 + u32::BITS);
        const PGE = 1 << (13 + u32::BITS);
        const MCA = 1 << (14 + u32::BITS);
        const CMOV = 1 << (15 + u32::BITS);
        const PAT = 1 << (16 + u32::BITS);
        const PSE36 = 1 << (17 + u32::BITS);
        const PSN = 1 << (18 + u32::BITS);
        const CLFSH = 1 << (19 + u32::BITS);
        const DS = 1 << (21 + u32::BITS);
        const ACPI = 1 << (22 + u32::BITS);
        const MMX = 1 << (23 + u32::BITS);
        const FXSR = 1 << (24 + u32::BITS);
        const SSE = 1 << (25 + u32::BITS);
        const SSE2 = 1 << (26 + u32::BITS);
        const SS = 1 << (27 + u32::BITS);
        const HTT = 1 << (28 + u32::BITS);
        const TM = 1 << (29 + u32::BITS);
        const PBE = 1 << (31 + u32::BITS);
    }
}

pub struct Leaf7(CpuidResult);

impl Leaf7 {
    pub fn get() -> Self {
        let result = cpuid(7, 0);
        Self(result)
    }

    pub fn subleaf1(&self) -> Option<Leaf7Subleaf1> {
        if self.max_subleaf() >= 1 {
            let result = cpuid(7, 1);
            Some(Leaf7Subleaf1(result))
        } else {
            None
        }
    }

    pub fn max_subleaf(&self) -> u32 {
        self.0.eax
    }
    pub fn flags(&self) -> Leaf7Flags {
        Leaf7Flags::from_bits_truncate(unsafe {
            core::mem::transmute::<[[u8; 4]; 4], u128>(
                [self.0.ebx, self.0.ecx, self.0.edx, 0].map(u32::to_ne_bytes),
            )
        })
    }
    pub fn mawau(&self) -> u8 {
        self.0.ecx.get_bits(17..22) as u8
    }
}

bitflags! {
    pub struct Leaf7Flags: u128{
        const FSGSBASE = 1 << 0;
        const IA32_TSC_ADJUST_MSR = 1 << 1;
        const SGX = 1 << 2;
        const BMI1 = 1 << 3;
        const HLE = 1 << 4;
        const AVX2 = 1 << 5;
        const SMEP = 1 << 7;
        const BMI2 = 1 << 8;
        const ERMS = 1 << 9;
        const INVPCID = 1 << 10;
        const RTM = 1 << 11;
        const PQM = 1 << 12;
        const FPU_CS_DS_DEPRECATION = 1 << 13;
        const MPX = 1 << 14;
        const PQE = 1 << 15;
        const AVX512F = 1 << 16;
        const AVX512DQ = 1 << 17;
        const RDSEED = 1 << 18;
        const ADX = 1 << 19;
        const SMAP = 1 << 20;
        const AVX512IFMA = 1 << 21;
        const PCOMMIT = 1 << 22;
        const CLFLUSHOPT = 1 << 23;
        const CLWB = 1 <<24 ;
        const INTEL_PT = 1 << 25;
        const AVX512PF = 1 << 26;
        const AVX512ER = 1 << 27;
        const AVX512CD = 1 << 28;
        const SHA = 1 << 29;
        const AVX512BW = 1 << 30;
        const AVX512VL = 1 << 31;

        const PREFETCHWT1 = 1 << (0 + u32::BITS);
        const AVX512VBMI = 1 << (1 + u32::BITS);
        const UMIP = 1 << (2 + u32::BITS);
        const PKU = 1 << (3 + u32::BITS);
        const OSPKE = 1 << (4 + u32::BITS);
        const RDPID = 1 << (22 + u32::BITS);
        const SGX_LC = 1 << (30 + u32::BITS);

        const SGX_KEYS = 1 << (1 + 2 * u32::BITS);
        const AVX512_4VNNIW = 1 << (2 + 2 * u32::BITS);
        const AVX512_4FMAPS = 1 << (3 + 2 * u32::BITS);
        const FSRM = 1 << (4 + 2 * u32::BITS);
        const UINTR = 1 << (5 + 2 * u32::BITS);
    }
}

pub struct Leaf7Subleaf1(CpuidResult);

impl Leaf7Subleaf1 {
    pub fn flags(&self) -> Leaf7Subleaf1Flags {
        Leaf7Subleaf1Flags::from_bits_truncate(unsafe {
            core::mem::transmute::<[u32; 4], u128>([self.0.eax, self.0.ebx, self.0.ecx, self.0.edx])
        })
    }
}

// in eax:ebx:ecx:edx order
bitflags! {
    pub struct Leaf7Subleaf1Flags: u128 {
        const SHA512 = 1 << 0;
        const SM3 = 1 << 1;
        const SM4 = 1 << 2;
        const RAO_INT = 1 << 3;
        const AVX_VNNI = 1 << 4;
        const AVX512_BF16 = 1 << 5;
        const LASS = 1 << 6;
        const CMPCCXADD = 1 << 7;
        const ARCHPERFMONTEXT = 1 << 8;
        const FZRM = 1 << 10;
        const FSRS = 1 << 11;
        const RSRCS = 1 << 12;
        const FRED = 1 << 17;
        const LKGS = 1 << 18;
        const WRMSRNS = 1 << 19;
        const NMI_SRC = 1 << 20;
        const AMX_FP16 = 1 << 21;
        const HRESET = 1 << 22;
        const AVX_IFMA = 1 << 23;
        const LAM = 1 << 26;
        const MSRLIST = 1 << 27;
        const INVD_DISABLE_POST_BIOS_DONE = 1 << 30;
        const MOVRS = 1 << 31;

        const PPINSTR = 1 << (0 + u32::BITS);
        const PBNDKB = 1 << (1 + u32::BITS);
        const CPUIDMAXVAL_LIM_RMV = 1 << (2 + u32::BITS);

        const RDT_M_ASYM = 1 << (0 + 2 * u32::BITS);
        const RDT_A_ASYM = 1 << (1 + 2 * u32::BITS);
        const MSR_IMM = 1 << (2 + 2 * u32::BITS);
        const ACE = 1 << (3 + 2 * u32::BITS);

        const AVX_VNNI_INT8 = 1 << (4 + 3 * u32::BITS);
        const AVX_NE_CONVERT = 1 << (5 + 3 * u32::BITS);
        const AMX_COMPLEX = 1 << (8 + 3 * u32::BITS);
        const AVX_VNNI_INT16 = 1 << (10 + 3 * u32::BITS);
        const UTMR = 1 << (13 + 3 * u32::BITS);
        const PREFETCHI = 1 << (14 + 3 * u32::BITS);
        const USER_MRS = 1 << (15 + 3 * u32::BITS);
        const UIRET_UIF_FROM_RFLAGS = 1 << (17 + 3 * u32::BITS);
        const CET_SSS = 1 << (18 + 3 * u32::BITS);
        const AVX10 = 1 << (19 + 3 * u32::BITS);
        const APX_F = 1 << (21 + 3 * u32::BITS);
        const SEC_TEE_ATTESTATION = 1 << (22 + 3 * u32::BITS);
        const MWAIT = 1 << (23 + 3 * u32::BITS);
        const SLSM = 1 << (24 + 3 * u32::BITS);
    }
}

bitflags! {
    pub struct Leaf8000001Flags: u64 {
        const FPU = 1 << 0;
        const VME = 1 << 1;
        const DE = 1 << 2;
        const PSE = 1 << 3;
        const TSC = 1 << 4;
        const MSR = 1 << 5;
        const PAE = 1 << 6;
        const MCE = 1 << 7;
        const CX8 = 1 << 8;
        const APIC = 1 << 9;
        const SYSCALL_K9 = 1 << 10;
        const SYSCALL = 1 << 11;
        const MTRR = 1 << 12;
        const PGE = 1 << 13;
        const MCA = 1 << 14;
        const CMOV = 1 << 15;
        const PAT = 1 << 16;
        const PSE36 = 1 << 17;
        const ECC_K9 = 1 << 18;
        const ECC = 1 << 19;
        const NX = 1 << 20;
        const MMXEXT = 1 << 22;
        const MMX = 1 << 23;
        const FXSR_OPT = 1 << 24;
        const PDPE1GB = 1 << 26;
        const RDTSCP = 1 << 27;
        const REX32 = 1 << 28;
        const LM = 1 << 29;
        const _3DNOWEXT = 1 << 30;
        const _3DNOW = 1 << 31;

        const LAHF_LM = 1 << (0 + u32::BITS);
        const CMP_LEGACY = 1 << (1 + u32::BITS);
        const SVM = 1 << (2 + u32::BITS);
        const EXTAPIC = 1 << (3 + u32::BITS);
        const CR8_LEGACY = 1 << (4 + u32::BITS);
        const LZCNT = 1 << (5 + u32::BITS);
        const SSE4A = 1 << (6 + u32::BITS);
        const MISALIGNSSE = 1 << (7 + u32::BITS);
        const _3DNOWPREFETCH = 1 << (8 + u32::BITS);
        const OSVW = 1 << (9 + u32::BITS);
        const IBS = 1 << (10 + u32::BITS);
        const XOP = 1 << (11 + u32::BITS);
        const SKINIT = 1 << (12 + u32::BITS);
        const WDT = 1 << (13 + u32::BITS);
        const LWP = 1 << (15 + u32::BITS);
        const FMA4 = 1 << (16 + u32::BITS);
        const TCE = 1 << (17 + u32::BITS);
        const NODEID_MSR = 1 << (19 + u32::BITS);
        const TBM = 1 << (21 + u32::BITS);
        const TOPOEXT = 1 << (22 + u32::BITS);
        const PERFCTR_CORE = 1 << (23 + u32::BITS);
        const PERFCTR_NB = 1 << (24 + u32::BITS);
        const DBX = 1 << (26 + u32::BITS);
        const PERFTSC = 1 << (27 + u32::BITS);
        const PCX_L2I = 1 << (28 + u32::BITS);
        const MONITORX = 1 << (29 + u32::BITS);
        const ADDR_MASK_EXT = 1 << (30 + u32::BITS);
    }
}

pub struct Leaf8000008(CpuidResult);

impl Leaf8000008 {
    pub fn get() -> Self {
        let result = cpuid(0x80000008, 0);
        Self(result)
    }
    pub fn physical_address_bits(&self) -> u8 {
        self.0.eax.get_bits(0..8) as u8
    }
    pub fn num_linear_address_bits(&self) -> u8 {
        self.0.eax.get_bits(8..16) as u8
    }
    pub fn guest_physical_address_bits(&self) -> u8 {
        self.0.eax.get_bits(16..24) as u8
    }

    pub fn num_physical_threads(&self) -> u8 {
        self.0.ecx.get_bits(0..8) as u8 + 1
    }
    pub fn apic_id_size(&self) -> u8 {
        self.0.ecx.get_bits(12..16) as u8
    }
    pub fn performance_timestamp_counter_size(&self) -> u8 {
        self.0.ecx.get_bits(16..18) as u8
    }
    pub fn max_invlpgb_page_count(&self) -> u16 {
        self.0.edx.get_bits(0..16) as u16
    }

    pub fn max_rdpru_ecx(&self) -> u16 {
        self.0.edx.get_bits(16..32) as u16
    }

    pub fn flags(&self) -> Leaf8000008Flags {
        Leaf8000008Flags::from_bits_truncate(self.0.ebx)
    }
}

impl Debug for Leaf8000008 {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("Leaf8000008")
            .field("physical_address_bits", &self.physical_address_bits())
            .field("num_linear_address_bits", &self.num_linear_address_bits())
            .field(
                "guest_physical_address_bits",
                &self.guest_physical_address_bits(),
            )
            .field("num_physical_threads", &self.num_physical_threads())
            .field("apic_id_size", &self.apic_id_size())
            .field(
                "performance_timestamp_counter_size",
                &self.performance_timestamp_counter_size(),
            )
            .field("max_invlpgb_page_count", &self.max_invlpgb_page_count())
            .field("max_rdpru_ecx", &self.max_rdpru_ecx())
            .field("flags", &self.flags())
            .finish()
    }
}

bitflags! {
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub struct Leaf8000008Flags: u32 {
        const CLZERO = 1 << 0;
        const RETIRED_INSTR = 1 << 1;
        const XRSTOR_FP_ERR = 1 << 2;
        const INVLPGB = 1 << 3;
        const RDPRU = 1 << 4;
        const XOTEXT = 1 << 5;
        const MBE = 1 << 6;
        const MCOMMIT = 1 << 8;
        const WBNOINVD = 1 << 9;
        const LBR_EXT_V1 = 1 << 10;
        const IBPB = 1 << 12;
        const WBINVD_INT = 1 << 13;
        const IBRS = 1 << 14;
        const STIBP = 1 << 15;
        const IBRS_ALWAYS = 1 << 16;
        const STIBP_ALWAYS = 1 << 17;
        const IBRS_PREFERRED = 1 << 18;
        const IBRS_SAME_MODE_PROT = 1 << 19;
        const NO_EFER_LMSLE = 1 << 20;
        const INVLPGB_NESTED = 1 << 21;
        const LBR_TSX = 1 << 22;
        const PPIN = 1 << 23;
        const SSBD = 1 << 24;
        const SSBD_LEGACY = 1 << 25;
        const SSBD_NO = 1 << 26;
        const CPPC = 1 << 27;
        const PSFD = 1 << 28;
        const BTC_NO = 1 << 29;
        const IBPB_RET = 1 << 30;
        const BRANCH_SAPLING = 1 << 31;
    }
}

impl CpuId {
    pub fn get() -> Self {
        let result = cpuid(0, 0);
        let vendor_id = unsafe {
            core::mem::transmute::<[u32; 3], [u8; 12]>([result.ebx, result.edx, result.ecx])
        };

        Self {
            max_eax: result.eax,
            vendor_id,
        }
    }

    pub fn vendor_id(&self) -> &[u8; 12] {
        &self.vendor_id
    }
}