#include <stdint.h>
#include <stdio.h>

#if defined(_MSC_VER)
#include <intrin.h>
#else
#include <cpuid.h>
#endif

enum {
    CPU_NO_X86_SIMD          = 0x000000,
    CPU_SSE2_SUPPORT         = 0x000001,
    CPU_SSE3_SUPPORT         = 0x000002,
    CPU_SSSE3_SUPPORT        = 0x000004,
    CPU_SSE4_1_SUPPORT       = 0x000008,
    CPU_SSE4_2_SUPPORT       = 0x000010,
    CPU_SSE4_A_SUPPORT       = 0x000020,
    CPU_FMA4_SUPPORT         = 0x000040,
    CPU_FMA3_SUPPORT         = 0x000080,
    CPU_AVX_SUPPORT          = 0x000100,
    CPU_AVX2_SUPPORT         = 0x000200,
    CPU_AVX512F_SUPPORT      = 0x000400,
    CPU_AVX512DQ_SUPPORT     = 0x000800,
    CPU_AVX512IFMA52_SUPPORT = 0x001000,
    CPU_AVX512PF_SUPPORT     = 0x002000,
    CPU_AVX512ER_SUPPORT     = 0x004000,
    CPU_AVX512CD_SUPPORT     = 0x008000,
    CPU_AVX512BW_SUPPORT     = 0x010000,
    CPU_AVX512VL_SUPPORT     = 0x020000,
    CPU_AVX512VBMI_SUPPORT   = 0x040000,
};

static inline void get_cpuid(int *array, int info_type)
{
#if defined(_MSC_VER)
    __cpuid(array, info_type);
#else
    __cpuid(info_type, array[0], array[1], array[2], array[3]);
#endif
}

static inline void get_cpuid2(int *array, int info_type, int ecx)
{
#if defined(_MSC_VER)
    __cpuidex(array, info_type, ecx);
#else
    __cpuid_count(info_type, ecx, array[0], array[1], array[2], array[3]);
#endif
}

static inline int is_bit_set(int bitfield, int bit)
{
    return bitfield & (1 << bit);
}

static uint32_t get_simd_support_info(void)
{
    uint32_t ret = 0;
    int regs[4];

    get_cpuid(regs, 0x00000001);
    if (is_bit_set(regs[3], 26)) {
        ret |= CPU_SSE2_SUPPORT;
    }
    if (is_bit_set(regs[2], 0)) {
        ret |= CPU_SSE3_SUPPORT;
    }
    if (is_bit_set(regs[2], 9)) {
        ret |= CPU_SSSE3_SUPPORT;
    }
    if (is_bit_set(regs[2], 19)) {
        ret |= CPU_SSE4_1_SUPPORT;
    }
    if (is_bit_set(regs[2], 26)) {
        ret |= CPU_SSE4_2_SUPPORT;
    }
    if (is_bit_set(regs[2], 27)) {
        if (is_bit_set(regs[2], 28)) {
            ret |= CPU_AVX_SUPPORT;
        }
        if (is_bit_set(regs[2], 12)) {
            ret |= CPU_FMA3_SUPPORT;
        }
    }

    get_cpuid(regs, 0x80000001);
    if (is_bit_set(regs[3], 6)) {
        ret |= CPU_SSE4_A_SUPPORT;
    }
    if (is_bit_set(regs[3], 16)) {
        ret |= CPU_FMA4_SUPPORT;
    }

    get_cpuid(regs, 0x00000000);
    if (regs[0] < 7) {
        return ret;
    }

    get_cpuid2(regs, 0x00000007, 0);
    if (is_bit_set(regs[1], 5)) {
        ret |= CPU_AVX2_SUPPORT;
    }
    if (!is_bit_set(regs[1], 16)) {
        return ret;
    }

    ret |= CPU_AVX512F_SUPPORT;
    if (is_bit_set(regs[1], 17)) {
        ret |= CPU_AVX512DQ_SUPPORT;
    }
    if (is_bit_set(regs[1], 21)) {
        ret |= CPU_AVX512IFMA52_SUPPORT;
    }
    if (is_bit_set(regs[1], 26)) {
        ret |= CPU_AVX512PF_SUPPORT;
    }
    if (is_bit_set(regs[1], 27)) {
        ret |= CPU_AVX512ER_SUPPORT;
    }
    if (is_bit_set(regs[1], 28)) {
        ret |= CPU_AVX512CD_SUPPORT;
    }
    if (is_bit_set(regs[1], 30)) {
        ret |= CPU_AVX512BW_SUPPORT;
    }
    if (is_bit_set(regs[1], 31)) {
        ret |= CPU_AVX512VL_SUPPORT;
    }
    if (is_bit_set(regs[2], 1)) {
        ret |= CPU_AVX512VBMI_SUPPORT;
    }

    return ret;
}

static void get_vendor_name(char *buff)
{
    int32_t regs[4];
    int32_t *temp = (int32_t *)buff;
    get_cpuid(regs, 0);
    temp[0] = regs[1];
    temp[1] = regs[3];
    temp[2] = regs[2];
}

static void get_brand_string(char *buff)
{
    int32_t regs[4];
    int32_t *temp = (int32_t *)buff;
    for (int i = 0; i < 3; i++) {
        get_cpuid(regs, 0x80000002 + i);
        for (int j = 0; j < 4; j++) {
            temp[i * 4 + j] = regs[j];
        }
    }
}

int main(void)
{
    char str[64] = {0};

    get_vendor_name(str);
    printf("vendor: %s\n", str);
    get_brand_string(str);
    printf("brand: %s\n", str);

    uint32_t cpu = get_simd_support_info();
    
    const struct {
        int flag;
        const char *name;
    } table[] = {
        {CPU_SSE2_SUPPORT        , "SSE2"      },
        {CPU_SSE3_SUPPORT        , "SSE3"      },
        {CPU_SSSE3_SUPPORT       , "SSSE3"     },
        {CPU_SSE4_1_SUPPORT      , "SSE4.1"    },
        {CPU_SSE4_2_SUPPORT      , "SSE4.2"    },
        {CPU_SSE4_A_SUPPORT      , "SSE4a"     },
        {CPU_FMA4_SUPPORT        , "FMA4"      },
        {CPU_FMA3_SUPPORT        , "FMA3"      },
        {CPU_AVX_SUPPORT         , "AVX"       },
        {CPU_AVX2_SUPPORT        , "AVX2"      },
        {CPU_AVX512F_SUPPORT     , "AVX512F"   },
        {CPU_AVX512DQ_SUPPORT    , "AVX512DQ"  },
        {CPU_AVX512IFMA52_SUPPORT, "AVX512IFMA"},
        {CPU_AVX512PF_SUPPORT    , "AVX512PF"  },
        {CPU_AVX512ER_SUPPORT    , "AVX512ER"  },
        {CPU_AVX512CD_SUPPORT    , "AVX512CD"  },
        {CPU_AVX512BW_SUPPORT    , "AVX512BW"  },
        {CPU_AVX512VL_SUPPORT    , "AVX512VL"  },
        {CPU_AVX512VBMI_SUPPORT  , "AVX512VBMI"},
        {-1, NULL}
    };

    if (cpu == CPU_NO_X86_SIMD) {
        printf("this cpu is too old.\n");
        return 0;
    }

    for (int i = 0; table[i].name != NULL; i++) {
        if (cpu & table[i].flag) {
            printf("%s is supported.\n", table[i].name);
        }
    }

    return 0;
}