#include "dispatch.h"
#include "cpu_detect.h"
#include <string.h>
#include <stdlib.h>

/* ============================================================================
 * External declarations for kernel implementations
 * ============================================================================ */

/* Scalar (always available) */
extern void convolve_scalar(float*, const float*, const float*,
                            unsigned int, unsigned int, unsigned int);
extern void convolve_transposed_scalar(float*, const float*, const float*,
                                       unsigned int, unsigned int);
extern void convolve_multi_filter_scalar(float*, const float*, const float*,
                                         unsigned int, unsigned int, unsigned int);
extern void upsample_polyphase_scalar(float*, const float*, const float*,
                                       unsigned int, unsigned int, unsigned int);

/* x86 implementations */
#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || defined(_M_IX86)

extern void convolve_sse2(float*, const float*, const float*,
                          unsigned int, unsigned int, unsigned int);
extern void convolve_transposed_sse2(float*, const float*, const float*,
                                     unsigned int, unsigned int);
extern void convolve_multi_filter_sse2(float*, const float*, const float*,
                                       unsigned int, unsigned int, unsigned int);
extern void upsample_polyphase_sse2(float*, const float*, const float*,
                                     unsigned int, unsigned int, unsigned int);

extern void convolve_sse4(float*, const float*, const float*,
                          unsigned int, unsigned int, unsigned int);
extern void convolve_transposed_sse4(float*, const float*, const float*,
                                     unsigned int, unsigned int);
extern void convolve_multi_filter_sse4(float*, const float*, const float*,
                                       unsigned int, unsigned int, unsigned int);

extern void convolve_avx(float*, const float*, const float*,
                         unsigned int, unsigned int, unsigned int);
extern void convolve_transposed_avx(float*, const float*, const float*,
                                    unsigned int, unsigned int);
extern void convolve_multi_filter_avx(float*, const float*, const float*,
                                      unsigned int, unsigned int, unsigned int);
extern void upsample_polyphase_avx(float*, const float*, const float*,
                                    unsigned int, unsigned int, unsigned int);

extern void convolve_avx2(float*, const float*, const float*,
                          unsigned int, unsigned int, unsigned int);
extern void convolve_transposed_avx2(float*, const float*, const float*,
                                     unsigned int, unsigned int);
extern void convolve_multi_filter_avx2(float*, const float*, const float*,
                                       unsigned int, unsigned int, unsigned int);
extern void upsample_polyphase_avx2(float*, const float*, const float*,
                                     unsigned int, unsigned int, unsigned int);

extern void convolve_avx512(float*, const float*, const float*,
                            unsigned int, unsigned int, unsigned int);
extern void convolve_transposed_avx512(float*, const float*, const float*,
                                       unsigned int, unsigned int);
extern void convolve_multi_filter_avx512(float*, const float*, const float*,
                                         unsigned int, unsigned int, unsigned int);

#endif /* x86 */

/* ARM implementations */
#if defined(__aarch64__) || defined(_M_ARM64)

extern void convolve_neon(float*, const float*, const float*,
                          unsigned int, unsigned int, unsigned int);
extern void convolve_transposed_neon(float*, const float*, const float*,
                                     unsigned int, unsigned int);
extern void convolve_multi_filter_neon(float*, const float*, const float*,
                                       unsigned int, unsigned int, unsigned int);
extern void upsample_polyphase_neon(float*, const float*, const float*,
                                     unsigned int, unsigned int, unsigned int);

#endif /* ARM */

/* ============================================================================
 * Dispatch table entry structure
 * ============================================================================ */

typedef struct {
    const char* name;
    unsigned int required_features;
    convolve_fn_t convolve;
    convolve_transposed_fn_t convolve_transposed;
    convolve_multi_filter_fn_t convolve_multi_filter;
    upsample_polyphase_fn_t upsample_polyphase;
} dispatch_entry_t;

/* ============================================================================
 * Dispatch table - ordered by preference (best first)
 * ============================================================================ */

static const dispatch_entry_t dispatch_table[] = {
#if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || defined(_M_IX86)
    {
        "AVX-512",
        CPU_FEATURE_AVX512F | CPU_FEATURE_AVX2 | CPU_FEATURE_FMA,
        convolve_avx512,
        convolve_transposed_avx512,
        convolve_multi_filter_avx512,
        upsample_polyphase_avx2  /* AVX2 polyphase is fine for AVX-512 */
    },
    {
        "AVX2+FMA",
        CPU_FEATURE_AVX2 | CPU_FEATURE_FMA,
        convolve_avx2,
        convolve_transposed_avx2,
        convolve_multi_filter_avx2,
        upsample_polyphase_avx2
    },
    {
        "AVX2",
        CPU_FEATURE_AVX2,
        convolve_avx2,
        convolve_transposed_avx2,
        convolve_multi_filter_avx2,
        upsample_polyphase_avx2
    },
    {
        "AVX",
        CPU_FEATURE_AVX,
        convolve_avx,
        convolve_transposed_avx,
        convolve_multi_filter_avx,
        upsample_polyphase_avx
    },
    {
        "SSE4.1",
        CPU_FEATURE_SSE41,
        convolve_sse4,
        convolve_transposed_sse4,
        convolve_multi_filter_sse4,
        upsample_polyphase_sse2
    },
    {
        "SSE2",
        CPU_FEATURE_SSE2,
        convolve_sse2,
        convolve_transposed_sse2,
        convolve_multi_filter_sse2,
        upsample_polyphase_sse2
    },
#endif /* x86 */

#if defined(__aarch64__) || defined(_M_ARM64)
    {
        "NEON",
        CPU_FEATURE_NEON,
        convolve_neon,
        convolve_transposed_neon,
        convolve_multi_filter_neon,
        upsample_polyphase_neon
    },
#endif /* ARM */

    /* Scalar fallback - always available, no required features */
    {
        "Scalar",
        0,
        convolve_scalar,
        convolve_transposed_scalar,
        convolve_multi_filter_scalar,
        upsample_polyphase_scalar
    },
};

#define NUM_DISPATCH_ENTRIES (sizeof(dispatch_table) / sizeof(dispatch_table[0]))

/* ============================================================================
 * Global state
 * ============================================================================ */

convolve_fn_t g_convolve = NULL;
convolve_transposed_fn_t g_convolve_transposed = NULL;
convolve_multi_filter_fn_t g_convolve_multi_filter = NULL;
upsample_polyphase_fn_t g_upsample_polyphase = NULL;
const char* g_impl_name = NULL;

static unsigned int g_cpu_features = 0;
static int g_initialized = 0;

/* ============================================================================
 * Dispatch functions
 * ============================================================================ */

void dispatch_init(void) {
    if (g_initialized) {
        return;
    }

    /* Check for environment variable override */
    const char* override = getenv("SGNL_CPU_IMPL");

    /* Detect CPU features */
    g_cpu_features = cpu_detect_features();

    /* If override specified, try to use it */
    if (override && override[0] != '\0') {
        if (dispatch_select(override) == 0) {
            g_initialized = 1;
            return;
        }
        /* Override failed, fall through to auto-detection */
    }

    /* Auto-select best available implementation */
    for (size_t i = 0; i < NUM_DISPATCH_ENTRIES; i++) {
        unsigned int required = dispatch_table[i].required_features;
        if ((g_cpu_features & required) == required) {
            g_convolve = dispatch_table[i].convolve;
            g_convolve_transposed = dispatch_table[i].convolve_transposed;
            g_convolve_multi_filter = dispatch_table[i].convolve_multi_filter;
            g_upsample_polyphase = dispatch_table[i].upsample_polyphase;
            g_impl_name = dispatch_table[i].name;
            g_initialized = 1;
            return;
        }
    }

    /* Should never reach here since Scalar has no requirements */
    g_convolve = convolve_scalar;
    g_convolve_transposed = convolve_transposed_scalar;
    g_convolve_multi_filter = convolve_multi_filter_scalar;
    g_upsample_polyphase = upsample_polyphase_scalar;
    g_impl_name = "Scalar";
    g_initialized = 1;
}

int dispatch_select(const char* name) {
    if (!name) {
        return -1;
    }

    /* Ensure features are detected */
    if (g_cpu_features == 0 && !g_initialized) {
        g_cpu_features = cpu_detect_features();
    }

    for (size_t i = 0; i < NUM_DISPATCH_ENTRIES; i++) {
        if (strcmp(dispatch_table[i].name, name) == 0) {
            /* Found matching entry, check if features available */
            unsigned int required = dispatch_table[i].required_features;
            if ((g_cpu_features & required) != required) {
                return -2;  /* Features not available */
            }

            g_convolve = dispatch_table[i].convolve;
            g_convolve_transposed = dispatch_table[i].convolve_transposed;
            g_convolve_multi_filter = dispatch_table[i].convolve_multi_filter;
            g_upsample_polyphase = dispatch_table[i].upsample_polyphase;
            g_impl_name = dispatch_table[i].name;
            return 0;  /* Success */
        }
    }

    return -1;  /* Not found */
}

const char** dispatch_get_available(void) {
    static const char* available[NUM_DISPATCH_ENTRIES + 1];
    static int cached = 0;

    if (cached) {
        return available;
    }

    /* Ensure features are detected */
    if (g_cpu_features == 0 && !g_initialized) {
        g_cpu_features = cpu_detect_features();
    }

    size_t count = 0;
    for (size_t i = 0; i < NUM_DISPATCH_ENTRIES; i++) {
        unsigned int required = dispatch_table[i].required_features;
        if ((g_cpu_features & required) == required) {
            available[count++] = dispatch_table[i].name;
        }
    }
    available[count] = NULL;
    cached = 1;

    return available;
}

unsigned int dispatch_get_cpu_features(void) {
    if (g_cpu_features == 0 && !g_initialized) {
        g_cpu_features = cpu_detect_features();
    }
    return g_cpu_features;
}
