#include "libblastrampoline_internal.h"
#include <complex.h>

/*
 * Some vendors (such as Accelerate) decided to trim off the trailing underscore
 * from the F77 symbol names when adding their ILP64 symbol names.  So the
 * symbol name `dgemm_` turns into `dgemm$NEWLAPACK$ILP64`.  But of course,
 * symbol names like `cblas_sdot` turn into `cblas_sdot$NEWLAPACK$ILP64`.
 *
 * So we need a way to selectively trim off the trailing underscore.  We do so
 * by shoving an ASCII "substitute character" onto the start of `$NEWLAPACK`,
 */
void build_symbol_name(char * out, const char *symbol_name, const char *suffix) {
    size_t symbol_len = strlen(symbol_name);
#if defined(SYMBOL_TRIMMING)
    if (suffix[0] == '\x1a') {
        if (symbol_name[symbol_len-1] == '_') {
            symbol_len -= 1;
        }
        suffix += 1;
    }
#endif

    strncpy(out, symbol_name, MAX_SYMBOL_LEN);
    strncpy(out + symbol_len, suffix, MAX_SYMBOL_LEN - symbol_len);
}

/*
 * Search for a symbol that ends in one of the given suffixes.  Returns NULL if not found.
 */
const char * symbol_suffix_search(void * handle, const char * symbol_name, const char ** suffixes, const int num_suffixes) {
    char symbol_name_suffixed[MAX_SYMBOL_LEN];
    for (int suffix_idx=0; suffix_idx<num_suffixes; suffix_idx++) {
        // Skip NULL suffixes (the case where we have no `suffix_hint`)
        if (suffixes[suffix_idx] == NULL) {
            continue;
        }

        build_symbol_name(symbol_name_suffixed, symbol_name, suffixes[suffix_idx]);
        if (lookup_symbol(handle, symbol_name_suffixed) != NULL) {
            return suffixes[suffix_idx];
        }
    }
    return NULL;
}

/*
 * Autodetect the name mangling suffix used by the given BLAS/LAPACK library.
 * We are primarily concerned with the FORTRAN names, so we automatically append
 * an underscore to the end, as is standard.
 */
const char * autodetect_symbol_suffix(void * handle, const char * suffix_hint) {
    // We auto-detect the symbol suffix of the given library by searching for common
    // BLAS and LAPACK symbols, combined with various suffixes that we know of.
    const char * symbol_names[] = {
        // fortran BLAS symbol
        "isamax_",
        // fortran LAPACK symbol
        "dpotrf_",
    };
    const char * suffixes[] = {
        // Possibly-NULL suffix that we should search over
        suffix_hint,

        // First, search for LP64-mangling suffixes, so that when we are loading libs from an
        // CLI environment, (where suffix hints are not easy) we want to give the most stable
        // configuration by default.
        "", "_", "__",

        // Next, ILP64-mangling suffixes
#if defined(_OS_DARWIN_) && defined(SYMBOL_TRIMMING)
        // Once again, search for Accelerate's non-pure-suffixed names
        "\x1a$NEWLAPACK$ILP64",
        // Apple Accelerate has an updated LAPACK interface.
        // Note that we are making use of our symbol trimming support here to eliminate
        // the F77 trailing underscore by starting the string with `\x1a`.
        "\x1a$NEWLAPACK",
#endif

        // Next, search for ILP64-mangling suffixes
        "64", "_64", "64_", "_64_", "_64__", "__64___",
    };

    // If the suffix hint is NULL, just skip it when calling `lookup_symbol()`.
    for (int symbol_idx=0; symbol_idx<sizeof(symbol_names)/sizeof(const char *); symbol_idx++) {
        const char * suffix = symbol_suffix_search(handle, symbol_names[symbol_idx], suffixes, sizeof(suffixes)/sizeof(const char *));
        if (suffix != NULL) {
            return suffix;
        }
    }
    return NULL;
}

/*
 * If this is a BLAS library, we'll check interface type by invoking `isamax` with a purposefully
 * incorrect `N` to cause it to change its return value based on how it is interpreting arugments.
 */
int32_t autodetect_blas_interface(void * isamax_addr) {
    if (env_lowercase_match("LBT_FORCE_INTERFACE", "ilp64")) {
        return LBT_INTERFACE_ILP64;
    }
    if (env_lowercase_match("LBT_FORCE_INTERFACE", "lp64")) {
        return LBT_INTERFACE_LP64;
    }
    // Typecast to function pointer for easier usage below
    int64_t (*isamax)(int64_t *, float *, int64_t *) = isamax_addr;

    // Purposefully incorrect length that is `3` on 32-bit, but massively negative on 64-bit
    int64_t n = 0xffffffff00000003;
    float X[3] = {1.0f, 2.0f, 1.0f};
    int64_t incx = 1;

    // Override `lsame_` to point to our `fake_lsame` if we're unable to `RTLD_DEEPBIND`
    if (use_deepbind == 0) {
        if (!push_fake_lsame()) {
            // Bad LBT compile?
            return LBT_INTERFACE_UNKNOWN;
        }
    }

    int64_t max_idx = isamax(&n, X, &incx);

    if (use_deepbind == 0) {
        if (!pop_fake_lsame()) {
            // Bad LBT compile?
            return LBT_INTERFACE_UNKNOWN;
        }
    }

    // Although we declare that `isamax` returns an `int64_t`, it may not actually do so,
    // since if it's an LP64 binary it's probably returning an `int32_t`.  Depending on the
    // register semantics, the high 32-bits may or may not be zeroed.  We don't really care
    // about them, since we know the two cases we're interested in.
    max_idx = max_idx & 0xffffffff;

    // This means the `isamax()` implementation saw `N < 0`, ergo it's a 64-bit library
    if (max_idx == 0) {
        return LBT_INTERFACE_ILP64;
    }
    // This means the `isamax()` implementation saw `N == 3`, ergo it's a 32-bit library
    if (max_idx == 2) {
        return LBT_INTERFACE_LP64;
    }
    // We have no idea what happened; `max_idx` isn't any of the options we thought it would be.
    return LBT_INTERFACE_UNKNOWN;
}

/*
 * If this is an LAPACK library, we'll check interface type by invoking `dpotrf` with a
 * purposefully incorrect `lda` to cause it to store an error code that we can inspect
 * and determine if the internal pointer dereferences were 32-bit or 64-bit.
 */
int32_t autodetect_lapack_interface(void * dpotrf_addr) {
    if (env_lowercase_match("LBT_FORCE_INTERFACE", "ilp64")) {
        return LBT_INTERFACE_ILP64;
    }
    if (env_lowercase_match("LBT_FORCE_INTERFACE", "lp64")) {
        return LBT_INTERFACE_LP64;
    }
    // Typecast to function pointer for easier usage below
    void (*dpotrf)(char *, int64_t *, double *, int64_t *, int64_t *) = dpotrf_addr;

    // This `dpotrf` invocation should result in an error code stored into `info`
    double testmat[4];
    char uplo = 'U';
    int64_t m = 2;
    int64_t lda = 0;
    int64_t info = 0;
    dpotrf(&uplo, &m, testmat, &lda, &info);

    if (info == 0xfffffffffffffffc) {
        // If `info` is actually `-4`, it means that `dpotrf` actually stored a 64-bit
        // value into `info`, which means that `lapack_int` internally is 64-bits
        return LBT_INTERFACE_ILP64;
    }
    if (info == 0x00000000fffffffc) {
        // This is what it looks like when a library stores a 32-bit value in a 64-bit slot.
        return LBT_INTERFACE_LP64;
    }
    // We have no idea what happened; `info` isn't any of the options we thought it would be.
    return LBT_INTERFACE_UNKNOWN;
}

/*
 * Autodetect the interface type of the given library with the given symbol mangling suffix.
 * Returns the values "32", "64" or "0", denoting the bitwidth of the internal index representation.
 */
int32_t autodetect_interface(void * handle, const char * suffix) {
    char symbol_name[MAX_SYMBOL_LEN];

    // Attempt BLAS `isamax()` test
    build_symbol_name(symbol_name, "isamax_", suffix);
    void * isamax = lookup_symbol(handle, symbol_name);
    if (isamax != NULL) {
        return autodetect_blas_interface(isamax);
    }

    // Attempt LAPACK `dpotrf()` test
    build_symbol_name(symbol_name, "dpotrf_", suffix);
    void * dpotrf = lookup_symbol(handle, symbol_name);
    if (dpotrf != NULL) {
        return autodetect_lapack_interface(dpotrf);
    }

    // Otherwise, this is probably not an LAPACK or BLAS library?!
    return LBT_INTERFACE_UNKNOWN;
}

int32_t autodetect_complex_return_style(void * handle, const char * suffix) {
    if (env_lowercase_match("LBT_FORCE_RETSTYLE", "normal")) {
        return LBT_COMPLEX_RETSTYLE_NORMAL;
    }
    if (env_lowercase_match("LBT_FORCE_RETSTYLE", "argument")) {
        return LBT_COMPLEX_RETSTYLE_ARGUMENT;
    }
    if (env_lowercase_match("LBT_FORCE_RETSTYLE", "fnda")) {
        return LBT_COMPLEX_RETSTYLE_FNDA;
    }

#ifdef COMPLEX_RETSTYLE_AUTODETECTION
    char symbol_name[MAX_SYMBOL_LEN];

    build_symbol_name(symbol_name, "zdotc_", suffix);
    void * zdotc_addr = lookup_symbol(handle, symbol_name);
    if (zdotc_addr == NULL) {
        return LBT_COMPLEX_RETSTYLE_UNKNOWN;
    }

    build_symbol_name(symbol_name, "cdotc_", suffix);
    void * cdotc_addr = lookup_symbol(handle, symbol_name);
    if (cdotc_addr == NULL) {
        return LBT_COMPLEX_RETSTYLE_UNKNOWN;
    }

    // Typecast to function pointer for easier usage below
    double complex (*zdotc_normal)(                  int64_t *, double complex *, int64_t *, double complex *, int64_t *) = zdotc_addr;
    void           (*zdotc_retarg)(double complex *, int64_t *, double complex *, int64_t *, double complex *, int64_t *) = zdotc_addr;

    // Typecast to function pointer for easier usage below
    float complex (*cdotc_normal)(                 int64_t *, float complex *, int64_t *, float complex *, int64_t *) = cdotc_addr;
    void          (*cdotc_retarg)(float complex *, int64_t *, float complex *, int64_t *, float complex *, int64_t *) = cdotc_addr;

    /*
     * First, check to see if `zdotc` zeros out the first argument if all arguments are zero.
     * Supposedly, most well-behaved implementations will return `0 + 0*I` if the length of
     * the inputs is zero; so if it is using a "return argument", that's a good way to find out.
     * 
     * We detect this by setting `retval` to an initial value of `-1` typecast to a complex
     * value.  The floating-point values are unimportant as they will be written to, but if
     * it is interpreted as an `int{32,64}_t`, it will be a negative value (which is not
     * allowed and should end the routine immediately).  This makes this invocation safe
     * across the full normal/argument, lp64/ilp64, cdotc/zdotc compatibility cube.
     */
    double complex retval_double = 0.0 + 1.0*I;
    int64_t zero = 0;
    double complex zeroc_double = 0.0 + 0.0*I;
    zdotc_retarg(&retval_double, &zero, &zeroc_double, &zero, &zeroc_double, &zero);

    /*
     * Next, do the same with `cdotc`, in order to detect situations where the ABI is
     * automatically inserting an extra argument to return 128-bit-wide values.
     * We call this `FNDA` for "Float Normal, Double Argument" style.
     */
    int64_t neg1 = -1;
    float complex retval_float = *(complex float *)(&neg1);
    float complex zeroc_float = 0.0f + 0.0f*I;
    cdotc_retarg(&retval_float, &zero, &zeroc_float, &zero, &zeroc_float, &zero);

    if (creal(retval_double) == 0.0 && cimag(retval_double) == 0.0) {
        // If the double values were reset, and the float values were also,
        // this is easy, we're just always argument-style:
        if (creal(retval_float) == 0.0f && cimag(retval_float) == 0.0f) {
            return LBT_COMPLEX_RETSTYLE_ARGUMENT;
        }

        // If the float values were not, let's try the normal return style:
        retval_float = 0.0f + 1.0f*I;
        retval_float = cdotc_normal(&zero, &zeroc_float, &zero, &zeroc_float, &zero);


        // If this works, we are in FNDA style (currently only observed on Windows x64)
        if (creal(retval_float) == 0.0f && cimag(retval_float) == 0.0f) {
            return LBT_COMPLEX_RETSTYLE_FNDA;
        }

        // Otherwise, cdotc is throwing a fit and we don't know what's up.
        return LBT_COMPLEX_RETSTYLE_UNKNOWN;
    }

    // If our double values were _not_ reset, let's hazard a guess that
    // we're dealing with a normal return style and test both types again:
    retval_double = 0.0 + 1.0*I;
    retval_double = zdotc_normal(&zero, &zeroc_double, &zero, &zeroc_double, &zero);
    retval_float = 0.0f + 1.0f*I;
    retval_float = cdotc_normal(&zero, &zeroc_float, &zero, &zeroc_float, &zero);


    // We only test for both working; we don't have a retstyle for float
    // being argument style and double being normal style.
    if ((creal(retval_double) == 0.0 && cimag(retval_double) == 0.0) &&
        (creal(retval_float) == 0.0f && cimag(retval_float) == 0.0f)) {
        return LBT_COMPLEX_RETSTYLE_NORMAL;
    }
#endif // COMPLEX_RETSTYLE_AUTODETECTION

    // If we get here, zdotc and cdotc are being uncooperative and we
    // do not appreciate it at all, not we don't my precious.
    return LBT_COMPLEX_RETSTYLE_UNKNOWN;
}

int32_t autodetect_f2c(void * handle, const char * suffix) {
    if (env_lowercase_match("LBT_FORCE_F2C", "plain")) {
        return LBT_F2C_PLAIN;
    }
    if (env_lowercase_match("LBT_FORCE_F2C", "required")) {
        return LBT_F2C_REQUIRED;
    }
#ifdef F2C_AUTODETECTION
    char symbol_name[MAX_SYMBOL_LEN];

    // Attempt BLAS `sdot()` test
    build_symbol_name(symbol_name, "sdot_", suffix);
    void * sdot_addr = lookup_symbol(handle, symbol_name);
    if (sdot_addr == NULL) {
        return LBT_F2C_UNKNOWN;
    }

    // Typecast to function pointer for easier usage below
    float  (*sdot_plain)(int64_t *, float *, int64_t *, float *, int64_t *) = sdot_addr;
    double (*sdot_f2c)(int64_t *, float *, int64_t *, float *, int64_t *)   = sdot_addr;

    // This `sdot_` invocation will return properly in either the `float` or `double` calling conventions
    float A[1] = {.5};
    float B[1] = {.5};
    int64_t n = 1;
    int64_t inca = 1;
    int64_t incb = 1;
    float plain = sdot_plain(&n, &A[0], &inca, &B[0], &incb);
    float f2c = sdot_f2c(&n, &A[0], &inca, &B[0], &incb);

    if (plain == 0.25) {
        // Normal calling convention
        return LBT_F2C_PLAIN;
    }
    if (f2c == 0.25) {
        // It's an f2c style calling convention
        return LBT_F2C_REQUIRED;
    }
#endif // F2C_AUTODETECTION

    // We have no idea what happened; nothing works and everything is broken
    return LBT_F2C_UNKNOWN;
}

int32_t autodetect_cblas_divergence(void * handle, const char * suffix) {
    if (env_lowercase_match("LBT_FORCE_CBLAS", "conformant")) {
        return LBT_CBLAS_CONFORMANT;
    }
    if (env_lowercase_match("LBT_FORCE_CBLAS", "divergent")) {
        return LBT_CBLAS_DIVERGENT;
    }

#ifdef CBLAS_DIVERGENCE_AUTODETECTION
    char symbol_name[MAX_SYMBOL_LEN];
    char extra_underscore_suffix[MAX_SYMBOL_LEN];
    snprintf(extra_underscore_suffix, MAX_SYMBOL_LEN, "_%s", suffix);

    build_symbol_name(symbol_name, "zdotc_", suffix);
    if (lookup_symbol(handle, symbol_name) != NULL ) {
        // If we have both `zdotc_64` and `cblas_zdotc_sub64`, it's all good:
        build_symbol_name(symbol_name, "cblas_zdotc_sub", suffix);
        if (lookup_symbol(handle, symbol_name) != NULL ) {
            return LBT_CBLAS_CONFORMANT;
        }

        // Do the fallback extra-underscore suffix search here, so we don't mistakenly
        // mark MKL v2024 as CBLAS-divergent
        build_symbol_name(symbol_name, "cblas_zdotc_sub", extra_underscore_suffix);
        if (lookup_symbol(handle, symbol_name) != NULL ) {
            return LBT_CBLAS_CONFORMANT;
        }

        const char * lp64_suffixes[] = {
            "", "_", "__",
        };
        const char * suffix = symbol_suffix_search(handle, "cblas_zdotc_sub", lp64_suffixes, sizeof(lp64_suffixes)/sizeof(const char *));

        // If we have `zdotc_64`, but we don't have `cblas_zdotc_sub64`, AND we have an
        // LP64-mangled `cblas_zdotc_sub`, then we know it's a library that has ILP64
        // FORTRAN symbols, but not CBLAS ones, such as MKL v2022.0.
        if (suffix != NULL) {
            return LBT_CBLAS_DIVERGENT;
        }
    }
#endif // CBLAS_DIVERGENCE_AUTODETECTION

    // If we can't even find `zdotc_64`, we don't know what this is.
    return LBT_CBLAS_UNKNOWN;
}
