#include "cpu_detect.h"
#include <string.h>
#include <stdio.h>

/* ============================================================================
 * x86/x86_64 CPU Detection
 * ============================================================================ */
#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || defined(_M_IX86)

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

/* Execute CPUID instruction */
static void cpuid(unsigned int leaf, unsigned int *eax, unsigned int *ebx,
                  unsigned int *ecx, unsigned int *edx) {
#ifdef _MSC_VER
    int info[4];
    __cpuid(info, (int)leaf);
    *eax = info[0];
    *ebx = info[1];
    *ecx = info[2];
    *edx = info[3];
#else
    __cpuid(leaf, *eax, *ebx, *ecx, *edx);
#endif
}

/* Execute CPUID with subleaf */
static void cpuid_count(unsigned int leaf, unsigned int subleaf,
                        unsigned int *eax, unsigned int *ebx,
                        unsigned int *ecx, unsigned int *edx) {
#ifdef _MSC_VER
    int info[4];
    __cpuidex(info, (int)leaf, (int)subleaf);
    *eax = info[0];
    *ebx = info[1];
    *ecx = info[2];
    *edx = info[3];
#else
    __cpuid_count(leaf, subleaf, *eax, *ebx, *ecx, *edx);
#endif
}

/* Read XCR0 via XGETBV to check OS support for extended registers */
static unsigned long long xgetbv(unsigned int index) {
#ifdef _MSC_VER
    return _xgetbv(index);
#else
    unsigned int eax, edx;
    __asm__ volatile("xgetbv" : "=a"(eax), "=d"(edx) : "c"(index));
    return ((unsigned long long)edx << 32) | eax;
#endif
}

unsigned int cpu_detect_features(void) {
    unsigned int features = 0;
    unsigned int eax, ebx, ecx, edx;
    unsigned int max_leaf;

    /* Get max supported CPUID leaf */
    cpuid(0, &max_leaf, &ebx, &ecx, &edx);
    if (max_leaf < 1) {
        return 0;
    }

    /* Leaf 1: Basic feature flags */
    cpuid(1, &eax, &ebx, &ecx, &edx);

    /* SSE2 (EDX bit 26) - baseline for x86_64 */
    if (edx & (1u << 26)) {
        features |= CPU_FEATURE_SSE2;
    }

    /* SSE3 (ECX bit 0) */
    if (ecx & (1u << 0)) {
        features |= CPU_FEATURE_SSE3;
    }

    /* SSSE3 (ECX bit 9) */
    if (ecx & (1u << 9)) {
        features |= CPU_FEATURE_SSSE3;
    }

    /* SSE4.1 (ECX bit 19) */
    if (ecx & (1u << 19)) {
        features |= CPU_FEATURE_SSE41;
    }

    /* SSE4.2 (ECX bit 20) */
    if (ecx & (1u << 20)) {
        features |= CPU_FEATURE_SSE42;
    }

    /* Check for AVX support: requires CPU support AND OS support */
    int has_avx = 0;
    if ((ecx & (1u << 28)) &&  /* AVX bit */
        (ecx & (1u << 27))) {   /* OSXSAVE bit - OS uses XSAVE */
        /* Verify OS has enabled XMM and YMM state saving */
        unsigned long long xcr0 = xgetbv(0);
        if ((xcr0 & 0x6) == 0x6) {  /* XMM (bit 1) + YMM (bit 2) */
            features |= CPU_FEATURE_AVX;
            has_avx = 1;
        }
    }

    /* FMA (ECX bit 12) - only useful with AVX */
    if (has_avx && (ecx & (1u << 12))) {
        features |= CPU_FEATURE_FMA;
    }

    /* Extended features (leaf 7) - AVX2, AVX-512 */
    if (has_avx && max_leaf >= 7) {
        cpuid_count(7, 0, &eax, &ebx, &ecx, &edx);

        /* AVX2 (EBX bit 5) */
        if (ebx & (1u << 5)) {
            features |= CPU_FEATURE_AVX2;
        }

        /* AVX-512 Foundation (EBX bit 16) */
        if (ebx & (1u << 16)) {
            /* AVX-512 requires OS support for ZMM registers */
            unsigned long long xcr0 = xgetbv(0);
            /* Check XMM, YMM, ZMM, and opmask are enabled */
            if ((xcr0 & 0xE6) == 0xE6) {
                features |= CPU_FEATURE_AVX512F;
            }
        }
    }

    return features;
}

/* ============================================================================
 * ARM/ARM64 CPU Detection
 * ============================================================================ */
#elif defined(__aarch64__) || defined(_M_ARM64)

#if defined(__linux__)
#include <sys/auxv.h>
#include <asm/hwcap.h>
#endif

unsigned int cpu_detect_features(void) {
    unsigned int features = 0;

#if defined(__APPLE__)
    /* All Apple Silicon has NEON and FMA */
    features |= CPU_FEATURE_NEON | CPU_FEATURE_NEON_FMA;

#elif defined(__linux__)
    unsigned long hwcap = getauxval(AT_HWCAP);

    /* ASIMD is NEON on AArch64 */
    if (hwcap & HWCAP_ASIMD) {
        features |= CPU_FEATURE_NEON;
    }

    /* FP is always present with ASIMD on AArch64, and includes FMA */
    if (hwcap & HWCAP_FP) {
        features |= CPU_FEATURE_NEON_FMA;
    }

#else
    /* Assume NEON is available on AArch64 (it's mandatory) */
    features |= CPU_FEATURE_NEON | CPU_FEATURE_NEON_FMA;
#endif

    return features;
}

/* ============================================================================
 * Generic / Unknown Architecture
 * ============================================================================ */
#else

unsigned int cpu_detect_features(void) {
    /* No SIMD features detected on unknown architecture */
    return 0;
}

#endif

/* ============================================================================
 * Feature String (all architectures)
 * ============================================================================ */
const char* cpu_features_string(unsigned int features) {
    static char buffer[256];
    buffer[0] = '\0';

    if (features & CPU_FEATURE_AVX512F) strcat(buffer, "AVX-512 ");
    if (features & CPU_FEATURE_AVX2)    strcat(buffer, "AVX2 ");
    if (features & CPU_FEATURE_FMA)     strcat(buffer, "FMA ");
    if (features & CPU_FEATURE_AVX)     strcat(buffer, "AVX ");
    if (features & CPU_FEATURE_SSE42)   strcat(buffer, "SSE4.2 ");
    if (features & CPU_FEATURE_SSE41)   strcat(buffer, "SSE4.1 ");
    if (features & CPU_FEATURE_SSSE3)   strcat(buffer, "SSSE3 ");
    if (features & CPU_FEATURE_SSE3)    strcat(buffer, "SSE3 ");
    if (features & CPU_FEATURE_SSE2)    strcat(buffer, "SSE2 ");
    if (features & CPU_FEATURE_NEON_FMA) strcat(buffer, "NEON+FMA ");
    else if (features & CPU_FEATURE_NEON) strcat(buffer, "NEON ");

    /* Remove trailing space */
    size_t len = strlen(buffer);
    if (len > 0 && buffer[len - 1] == ' ') {
        buffer[len - 1] = '\0';
    }

    if (buffer[0] == '\0') {
        strcpy(buffer, "Scalar");
    }

    return buffer;
}
