# Exact token-id parity for the Qwen2 byte-level BPE tokenizer against
# HuggingFace reference cases (generated by
# tools/gen_qwen_tokenizer_cases.py, checked in). The 11 MB
# tokenizer.json is not shipped; the test locates one and skips if
# absent.

library(diffuseR)

find_qwen_tokenizer <- function() {
  p <- Sys.getenv("DIFFUSER_QWEN_TOKENIZER", "")
  if (nzchar(p) && file.exists(p)) {
    return(p)
  }
  if (requireNamespace("hfhub", quietly = TRUE)) {
    p <- tryCatch(
      suppressMessages(hfhub::hub_download(
        "black-forest-labs/FLUX.2-klein-4B",
        "tokenizer/tokenizer.json", local_files_only = TRUE
      )),
      error = function(e) ""
    )
    if (nzchar(p) && file.exists(p)) {
      return(p)
    }
  }
  p <- "../../tools/cache/tokenizer_qwen.json"
  if (file.exists(p)) {
    return(p)
  }
  ""
}

tok_path <- find_qwen_tokenizer()
if (!nzchar(tok_path)) exit_file("no Qwen tokenizer.json available")

cases_path <- system.file("tinytest", "fixtures", "qwen_tokenizer_cases.json",
  package = "diffuseR")
if (cases_path == "") cases_path <- "fixtures/qwen_tokenizer_cases.json"
if (!file.exists(cases_path)) exit_file("qwen tokenizer cases missing")

cases <- jsonlite::fromJSON(cases_path, simplifyVector = FALSE)

tok <- qwen_bpe_tokenizer(tok_path)
expect_equal(tok$pad_id, 151643L)

# --- raw encodings (no template, no padding) --------------------------------------

for (case in cases$cases) {
  expected <- as.integer(unlist(case$ids))
  got <- encode_qwen(tok, case$text, max_length = NULL,
    chat_template = FALSE)$input_ids[[1]]
  expect_equal(got, expected,
    info = sprintf("text: %s", substr(case$text, 1, 60)))
}

# --- chat template + padding/truncation ----------------------------------------------

for (case in cases$templated) {
  got <- encode_qwen(tok, case$text, max_length = case$max_length,
    chat_template = TRUE)
  expect_equal(as.integer(got$input_ids[1, ]), as.integer(unlist(case$ids)),
    info = sprintf("templated: %s", substr(case$text, 1, 60)))
  expect_equal(as.integer(got$attention_mask[1, ]),
    as.integer(unlist(case$mask)))
}

# --- batch shape ------------------------------------------------------------------------

batch <- encode_qwen(tok, c("a photo of a cat", "hello"), max_length = 32L)
expect_equal(dim(batch$input_ids), c(2L, 32L))
expect_equal(dim(batch$attention_mask), c(2L, 32L))
