#pragma once

#include <algorithm>
#include <cpuid.h>
#include <vector>

#ifdef _MSC_VER
#include <intrin.h>
#else
#include <cpuid.h>
#endif

inline void get_cpuid(int info[4], int infoType) {
#ifdef _MSC_VER
    __cpuidex(info, infoType, 0);
#else
    __cpuid_count(infoType, 0, info[0], info[1], info[2], info[3]);
#endif
}

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

inline std::vector<simd_type> get_simd_support_type() {
    unsigned int eax, ebx, ecx, edx;
    __get_cpuid(1, &eax, &ebx, &ecx, &edx);
    std::vector<simd_type> simd_types;
#if CPU_AMD64
    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) {
        simd_types.push_back(simd_type::avx512);
    }
#endif
#if CPU_ARM
    uint64_t id_aa64isar0_ = 0;
    uint64_t id_aa64pfr0_ = 0;

    // Reading the ID_AA64ISAR0_EL1 register
    asm volatile("mrs %0, ID_AA64ISAR0_EL1" : "=r" (id_aa64isar0_));
    // Reading the ID_AA64PFR0_EL1 register
    asm volatile("mrs %0, ID_AA64PFR0_EL1" : "=r" (id_aa64pfr0_));
    if ((id_aa64isar0_ >> 24) & 0xf) {
        simd_types.push_back(simd_type::neon64);
    }
    if ((id_aa64isar0_ >> 28) & 0xf) {
        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_neon64() || support_neon128();
    }

    [[nodiscard]] bool support_neon64() const {
        return support_simd(simd_type::neon64);
    }

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