#pragma once

#include <algorithm>
#include <vector>

#ifdef _MSC_VER
#include <intrin.h>
#else
#if CPU_AMD64
#include <cpuid.h>
#else
#include <sys/sysctl.h>
#endif
#endif

enum class simd_type {
    sse2,
    sse3,
    ssse3,
    sse41,
    sse42,
    avx,
    avx2,
    avx512,
    neon,
};

inline std::vector<simd_type> get_simd_support_type() {
    std::vector<simd_type> simd_types;
#if CPU_AMD64
    unsigned int eax, ebx, ecx, edx;
    __get_cpuid(1, &eax, &ebx, &ecx, &edx);
    if (ecx & bit_SSE2) {
        simd_types.push_back(simd_type::sse2);
    }
    if (ecx & bit_SSE3) {
        simd_types.push_back(simd_type::sse3);
    }
    if (ecx & bit_SSSE3) {
        simd_types.push_back(simd_type::ssse3);
    }
    if (ecx & bit_SSE4_1) {
        simd_types.push_back(simd_type::sse41);
    }
    if (ecx & bit_SSE4_2) {
        simd_types.push_back(simd_type::sse42);
    }
    if (ecx & bit_AVX) {
        simd_types.push_back(simd_type::avx);
    }
    if (ebx & bit_AVX2) {
        simd_types.push_back(simd_type::avx2);
    }
    if (ebx & bit_AVX512F || ebx & bit_AVX512DQ || ebx & bit_AVX512VL || ebx & bit_AVX512BW ||
        ebx & bit_AVX512PF || ebx & bit_AVX512IFMA || ebx & bit_AVX512CD || ebx & bit_AVX512ER) {
        simd_types.push_back(simd_type::avx512);
    }
#endif
#if CPU_ARM
    int neon_support = 0;
    size_t size = sizeof(neon_support);
    if (sysctlbyname("hw.optional.neon", &neon_support, &size, NULL, 0) == 0) {
        if (neon_support) {
            simd_types.push_back(simd_type::neon64);
            simd_types.push_back(simd_type::neon128);
        }
    }
#endif

    return simd_types;
}

class CORE_API cpuid {
public:
    cpuid() {
        simd_types = get_simd_support_type();
    }
    [[nodiscard]] bool support_simd(simd_type simd) const {
        return std::ranges::find(simd_types, simd) != simd_types.end();
    }

    [[nodiscard]] bool support_sse() const {
        return support_simd(simd_type::sse42) || support_simd(simd_type::sse41) || support_simd(simd_type::ssse3) || support_simd(simd_type::sse3) || support_simd(simd_type::sse2);
    }

    [[nodiscard]] bool support_avx() const {
        return support_simd(simd_type::avx);
    }

    [[nodiscard]] bool support_avx2() const {
        return support_simd(simd_type::avx2);
    }

    [[nodiscard]] bool support_avx512() const {
        return support_simd(simd_type::avx512);
    }

    [[nodiscard]] bool support_neon() const {
        return support_simd(simd_type::neon);
    }
private:
    std::vector<simd_type> simd_types;
};