// qwen3_asr_bpe_parity.cpp - byte-level BPE encoder parity test.
//
// Loads the Qwen3-ASR-0.6B GGUF's tokenizer and runs two batteries
// of test cases:
//
//   1. k_bpe_fixtures — plain-text strings (no Qwen3-ASR special
//      tokens). Tokenizer::encode() must produce token-for-token
//      the same ids as Hugging Face's reference tokenizer would for
//      encode(add_special_tokens=False).
//
//   2. k_lang_prefix_fixtures — the 30 canonical publisher language
//      names. encode_language_prefix() (the qwen3_asr internal that
//      assembles "language {Name}<asr_text>" for the chat template)
//      must match HF's full-prefix tokenization, including the
//      <asr_text> special-token id it appends by hand.
//
// The fixture data is generated by
// scripts/tokenizer-parity-fixture.py from the real HF tokenizer
// bundled with the 0.6B checkpoint. If Qwen3-ASR ships a vocab
// update the script is re-run and this file is unchanged.
//
// Gated on the TRANSCRIBE_QWEN3_ASR_0_6B_GGUF env var (same var the
// structural real-smoke test uses). Unset -> exit 77 (skip).

#include "arch/qwen3_asr/qwen3_asr.h"
#include "transcribe-model.h"
#include "transcribe-tokenizer.h"
#include "transcribe.h"

#include <sys/stat.h>

#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>

// encode_language_prefix() is declared in arch/qwen3_asr/qwen3_asr.h
// (which this test already pulls in) so no further wiring is needed.
// The parity test uses it directly.

namespace {

int g_failures = 0;

// Fixture schema used by the generated .inc file.
struct Fixture {
    const char *  text;
    const int32_t ids[32];
    size_t        n_ids;
};

#include "fixtures/qwen3_asr_bpe_parity.inc"

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

void check_ids_equal(const char *                 label,
                     const char *                 text,
                     const int32_t *              expected,
                     size_t                       expected_n,
                     const std::vector<int32_t> & actual) {
    if (actual.size() == expected_n && std::memcmp(actual.data(), expected, expected_n * sizeof(int32_t)) == 0) {
        return;
    }

    ++g_failures;
    std::fprintf(stderr, "FAIL[%s] input=\"%s\"\n  expected (%zu):", label, text, expected_n);
    for (size_t i = 0; i < expected_n; ++i) {
        std::fprintf(stderr, " %d", static_cast<int>(expected[i]));
    }
    std::fprintf(stderr, "\n  actual   (%zu):", actual.size());
    for (int32_t id : actual) {
        std::fprintf(stderr, " %d", static_cast<int>(id));
    }
    std::fprintf(stderr, "\n");
}

}  // namespace

int main() {
    const char * env = std::getenv("TRANSCRIBE_QWEN3_ASR_0_6B_GGUF");
    if (env == nullptr || env[0] == '\0') {
        std::fprintf(stderr,
                     "qwen3_asr_bpe_parity: TRANSCRIBE_QWEN3_ASR_0_6B_GGUF "
                     "not set; skipping.\n");
        return 77;
    }
    const std::string model_path = env;
    if (!file_exists(model_path)) {
        std::fprintf(stderr, "qwen3_asr_bpe_parity: model not found: %s\n", model_path.c_str());
        return 77;
    }

    transcribe_model_load_params mp;
    transcribe_model_load_params_init(&mp);
    mp.backend                      = TRANSCRIBE_BACKEND_CPU;
    struct transcribe_model * model = nullptr;
    if (transcribe_model_load_file(model_path.c_str(), &mp, &model) != TRANSCRIBE_OK || model == nullptr) {
        std::fprintf(stderr, "qwen3_asr_bpe_parity: failed to load model\n");
        return EXIT_FAILURE;
    }

    const auto *                  base = reinterpret_cast<const transcribe_model *>(model);
    const auto *                  qm   = static_cast<const transcribe::qwen3_asr::QwenAsrModel *>(base);
    const transcribe::Tokenizer & tok  = qm->tok;

    if (!tok.has_encoder()) {
        std::fprintf(stderr,
                     "qwen3_asr_bpe_parity: tokenizer has no encoder "
                     "(model=\"%s\", merges loaded? no)\n",
                     tok.model_type().c_str());
        transcribe_model_free(model);
        return EXIT_FAILURE;
    }

    // Section 1: plain-text BPE parity.
    std::fprintf(stderr, "qwen3_asr_bpe_parity: %zu plain BPE fixtures\n", k_bpe_fixtures_n);
    for (size_t i = 0; i < k_bpe_fixtures_n; ++i) {
        const auto &            f = k_bpe_fixtures[i];
        std::vector<int32_t>    got;
        const transcribe_status st = tok.encode(f.text, got);
        if (st != TRANSCRIBE_OK) {
            std::fprintf(stderr, "FAIL[bpe] encode returned %s for input \"%s\"\n", transcribe_status_string(st),
                         f.text);
            ++g_failures;
            continue;
        }
        check_ids_equal("bpe", f.text, f.ids, f.n_ids, got);
    }

    // Section 2: language-prefix parity (encode_language_prefix).
    std::fprintf(stderr, "qwen3_asr_bpe_parity: %zu language-prefix fixtures\n", k_lang_prefix_fixtures_n);
    for (size_t i = 0; i < k_lang_prefix_fixtures_n; ++i) {
        const auto &            f = k_lang_prefix_fixtures[i];
        std::vector<int32_t>    got;
        const transcribe_status st = transcribe::qwen3_asr::encode_language_prefix(tok, f.bcp47, got);
        if (st != TRANSCRIBE_OK) {
            std::fprintf(stderr,
                         "FAIL[lang] encode_language_prefix returned "
                         "%s for bcp47=\"%s\"\n",
                         transcribe_status_string(st), f.bcp47);
            ++g_failures;
            continue;
        }
        char label[64];
        std::snprintf(label, sizeof(label), "lang %s", f.bcp47);
        check_ids_equal(label, f.pub_name, f.ids, f.n_ids, got);
    }

    transcribe_model_free(model);

    if (g_failures > 0) {
        std::fprintf(stderr, "qwen3_asr_bpe_parity: %d failures\n", g_failures);
        return EXIT_FAILURE;
    }
    std::fprintf(stdout, "qwen3_asr_bpe_parity: ok\n");
    return EXIT_SUCCESS;
}