// Loader gate for the dspark GIDD log-SNR metadata contract
// (src/models/dspark.cpp load_arch_hparams):
//
//   * dspark.log_snr_conditioning = true makes min/max_log_snr REQUIRED
//     (a silently-defaulted bound would collapse (max_snr - min_snr) to zero
//     in the featurizer and fill the conditioning input with NaNs),
//   * both bounds must be finite,
//   * max_log_snr must be strictly greater than min_log_snr,
//   * with conditioning off/absent, missing bounds stay allowed (every
//     pre-GIDD drafter ships without the keys).
//
// Each case is a tiny random-weight GGUF variant; the two variants that are
// expected to LOAD assert full llama_model_load_from_file success, and every
// rejected variant additionally asserts the captured llama_log output
// mentions log_snr/log-SNR, so a pass means THE GUARD fired rather than some
// unrelated load failure.
//
// Needs the variant suite generated by the companion script; the test skips
// (exit 0) when the directory holds none of the variants:
//   python3 tests/gen-tiny-dspark.py --suite /tmp/dspark-logsnr-variants
//   test-dspark-logsnr-meta /tmp/dspark-logsnr-variants
//
// usage: test-dspark-logsnr-meta <variants-dir>

#include "ggml-backend.h"
#include "llama.h"

#include <sys/stat.h>

#include <cstdio>
#include <string>

// keep in sync with SUITE in tests/gen-tiny-dspark.py
struct variant_case {
    const char * file;
    const char * desc;
    bool         expect_load;
};

static const variant_case EXPECTED_VARIANTS[] = {
    { "logsnr-valid.gguf",    "conditioning=1 min=-9 max=6 (finite, ordered)", true  },
    { "logsnr-omit-min.gguf", "conditioning=1 min omitted",                    false },
    { "logsnr-omit-max.gguf", "conditioning=1 max omitted",                    false },
    { "logsnr-min-nan.gguf",  "conditioning=1 min=nan",                        false },
    { "logsnr-max-inf.gguf",  "conditioning=1 max=inf",                        false },
    { "logsnr-equal.gguf",    "conditioning=1 min==max",                       false },
    { "logsnr-reversed.gguf", "conditioning=1 min>max",                        false },
    { "logsnr-cond-off.gguf", "conditioning=0 bounds omitted",                 true  },
};

static std::string g_log;  // accumulated llama log output for the current load

static void log_capture(ggml_log_level level, const char * text, void * /*user_data*/) {
    GGML_UNUSED(level);
    g_log += text;
}

static bool file_exists(const std::string & path) {
    struct stat st;
    return stat(path.c_str(), &st) == 0;
}

int main(int argc, char ** argv) {
    if (argc < 2) {
        fprintf(stderr, "usage: %s <variants-dir>\n", argv[0]);
        return 1;
    }
    const std::string dir = argv[1];

    // skip-if-absent (mirrors the gen-tiny-qwen35.py-backed tests): the suite
    // is generated manually, so a missing directory is not a failure. a
    // PARTIAL suite, however, is a generation error and must not skip.
    size_t n_present = 0;
    for (const auto & v : EXPECTED_VARIANTS) {
        n_present += file_exists(dir + "/" + v.file) ? 1 : 0;
    }
    if (n_present == 0) {
        fprintf(stderr, "%s : skipping, no variants in '%s'\n", __func__, dir.c_str());
        fprintf(stderr, "%s : generate them with: python3 tests/gen-tiny-dspark.py --suite '%s'\n", __func__,
                dir.c_str());
        return 0;
    }
    if (n_present != sizeof(EXPECTED_VARIANTS) / sizeof(EXPECTED_VARIANTS[0])) {
        fprintf(stderr, "%s : FAIL: '%s' holds a partial variant suite (%zu present), regenerate it\n", __func__,
                dir.c_str(), n_present);
        return 1;
    }

    llama_log_set(log_capture, nullptr);
    ggml_backend_load_all();  // no-op unless built with GGML_BACKEND_DL
    llama_backend_init();

    llama_model_params mparams = llama_model_default_params();
    mparams.n_gpu_layers       = 0;  // CPU is all this needs; keep it backend-independent

    int n_failed = 0;

    for (const auto & v : EXPECTED_VARIANTS) {
        const std::string path = dir + "/" + v.file;

        g_log.clear();
        llama_model *     model = llama_model_load_from_file(path.c_str(), mparams);
        const std::string log   = g_log;

        bool ok;
        if (v.expect_load) {
            ok = model != nullptr;
            if (!ok) {
                fprintf(stderr, "  load log:\n%s\n", log.c_str());
                fprintf(stderr, "FAIL %-22s (%s): expected load success, got failure\n", v.file, v.desc);
            }
        } else {
            const bool guard_fired =
                log.find("log_snr") != std::string::npos || log.find("log-SNR") != std::string::npos;
            ok = model == nullptr && guard_fired;
            if (model != nullptr) {
                fprintf(stderr, "FAIL %-22s (%s): expected load failure, model loaded\n", v.file, v.desc);
            } else if (!guard_fired) {
                fprintf(stderr, "  load log:\n%s\n", log.c_str());
                fprintf(stderr, "FAIL %-22s (%s): load failed but not on the log-SNR guard\n", v.file, v.desc);
            }
        }

        if (model != nullptr) {
            llama_model_free(model);
        }

        if (ok) {
            fprintf(stderr, "OK   %-22s (%s): %s\n", v.file, v.desc,
                    v.expect_load ? "loaded" : "rejected by the log-SNR guard");
        } else {
            n_failed++;
        }
    }

    llama_backend_free();

    if (n_failed > 0) {
        fprintf(stderr, "%s : %d variant(s) FAILED\n", __func__, n_failed);
        return 1;
    }

    fprintf(stderr, "%s : all %zu log-SNR metadata variants behaved as expected\n", __func__,
            sizeof(EXPECTED_VARIANTS) / sizeof(EXPECTED_VARIANTS[0]));
    return 0;
}