#include <stdio.h>
#include <cpuid.h>

typedef struct {
    unsigned int bit;
    const char *name;
    const char *desc;
} CpuFeature;

CpuFeature leaf1_edx[] = {
    {1 << 0,  "FPU",    "Onboard x87 Floating-Point Unit"},
    {1 << 4,  "TSC",    "Time Stamp Counter"},
    {1 << 9,  "APIC",   "On-Chip APIC Hardware"},
    {1 << 15, "CMOV",   "Conditional Move Instructions"},
    {1 << 23, "MMX",    "MMX Technology"},
    {1 << 25, "SSE",    "Streaming SIMD Extensions"},
    {1 << 26, "SSE2",   "Streaming SIMD Extensions 2"},
    {1 << 28, "HTT",    "Max APIC IDs / Hyper-Threading"},
    {0, NULL, NULL}
};

CpuFeature leaf1_ecx[] = {
    {1 << 0,  "SSE3",   "Streaming SIMD Extensions 3"},
    {1 << 9,  "SSSE3",  "Supplemental Streaming SIMD Extensions 3"},
    {1 << 12, "FMA",    "Fused Multiply-Add (FMA3)"},
    {1 << 19, "SSE4_1", "Streaming SIMD Extensions 4.1"},
    {1 << 20, "SSE4_2", "Streaming SIMD Extensions 4.2"},
    {1 << 25, "AES",    "AES New Instructions"},
    {1 << 26, "XSAVE",  "XSAVE/XRSTOR States"},
    {1 << 27, "OSXSAVE","OS Enabled XSAVE/XRSTOR"},
    {1 << 28, "AVX",    "Advanced Vector Extensions"},
    {1 << 29, "F16C",   "16-bit Floating-Point Conversion"},
    {1 << 30, "RDRAND", "RDRAND Instruction"},
    {0, NULL, NULL}
};

CpuFeature leaf7_ebx[] = {
    {1 << 3,  "BMI1",   "Bit Manipulation Instruction Set 1"},
    {1 << 5,  "AVX2",   "Advanced Vector Extensions 2"},
    {1 << 8,  "BMI2",   "Bit Manipulation Instruction Set 2"},
    {1 << 16, "AVX512F","AVX-512 Foundation"},
    {1 << 17, "AVX512DQ","AVX-512 Doubleword and Quadword"},
    {1 << 18, "RDSEED", "RDSEED Instruction"},
    {1 << 19, "ADX",    "Intel ADX (Multi-Precision Add-Carry)"},
    {1 << 29, "SHA",    "Intel SHA Extensions"},
    {1 << 30, "AVX512BW","AVX-512 Byte and Word"},
    {0, NULL, NULL}
};

CpuFeature leaf7_ecx[] = {
    {1 << 1,  "AVX512VBMI", "AVX-512 Vector Byte Manipulation"},
    {1 << 11, "AVX512VNNI", "AVX-512 Vector Neural Network Instructions"},
    {1 << 12, "AVX512BITALG","AVX-512 Bit Algorithms"},
    {0, NULL, NULL}
};

void print_raw_registers(unsigned int leaf, unsigned int subleaf, unsigned int eax, unsigned int ebx, unsigned int ecx, unsigned int edx) {
    printf("\n=========================================================\n");
    printf("CPUID Leaf 0x%X (Sub-leaf 0x%X) Raw Hex Output:\n", leaf, subleaf);
    printf("  EAX: 0x%08X    EBX: 0x%08X    ECX: 0x%08X    EDX: 0x%08X\n", eax, ebx, ecx, edx);
    printf("=========================================================\n");
}

void print_features(const char *heading, unsigned int reg, CpuFeature *features) {
    printf("\n[%s Decoded Flags]\n", heading);
    printf("%-12s | %-8s | %s\n", "Status", "Feature", "Description");
    printf("---------------------------------------------------------\n");
    for (int i = 0; features[i].name != NULL; i++) {
        int supported = (reg & features[i].bit) != 0;
        printf("%-12s | %-8s | %s\n", 
               supported ? "  [ YES ]" : "  [  NO ]", 
               features[i].name, 
               features[i].desc);
    }
}

int main(void) {
    unsigned int eax, ebx, ecx, edx;
    unsigned int max_leaf;

    if (!__get_cpuid(0, &max_leaf, &ebx, &ecx, &edx)) {
        printf("Error: CPUID instruction not supported.\n");
        return 1;
    }
    
    printf("--- Checking Hardware Capabilities via CPUID ---\n");
    print_raw_registers(0, 0, max_leaf, ebx, ecx, edx);

    // Parse Leaf 1
    if (max_leaf >= 1) {
        __get_cpuid(1, &eax, &ebx, &ecx, &edx);
        print_raw_registers(1, 0, eax, ebx, ecx, edx);
        print_features("Leaf 1 EDX", edx, leaf1_edx);
        print_features("Leaf 1 ECX", ecx, leaf1_ecx);
    }

    // Parse Leaf 7
    if (max_leaf >= 7) {
        __get_cpuid_count(7, 0, &eax, &ebx, &ecx, &edx);
        print_raw_registers(7, 0, eax, ebx, ecx, edx);
        print_features("Leaf 7 EBX", ebx, leaf7_ebx);
        print_features("Leaf 7 ECX", ecx, leaf7_ecx);
    }

    return 0;
}