# Exact token-id parity for the Unigram (T5/SentencePiece) tokenizer
# against HuggingFace tokenizers reference cases (generated by
# tools/gen_t5_tokenizer_cases.py, checked in).
#
# The tokenizer.json itself is not shipped (2.4 MB, and FLUX's copy is
# behind the gated repo); the test locates one and skips if absent.

library(diffuseR)

find_t5_tokenizer <- function() {
  # 1. Explicit override
  p <- Sys.getenv("DIFFUSER_T5_TOKENIZER", "")
  if (nzchar(p) && file.exists(p)) {
    return(p)
  }
  # 2. FLUX.1-schnell HF cache (post-download_flux1)
  if (requireNamespace("hfhub", quietly = TRUE)) {
    p <- tryCatch(
      suppressMessages(hfhub::hub_download(
        "black-forest-labs/FLUX.1-schnell",
        "tokenizer_2/tokenizer.json", local_files_only = TRUE
      )),
      error = function(e) ""
    )
    if (nzchar(p) && file.exists(p)) {
      return(p)
    }
  }
  # 3. Dev copy in the source tree (tools/gen_t5_tokenizer_cases.py)
  p <- "../../tools/cache/tokenizer_t5.json"
  if (file.exists(p)) {
    return(p)
  }
  ""
}

tok_path <- find_t5_tokenizer()
if (!nzchar(tok_path)) exit_file("no T5 tokenizer.json available")

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

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

tok <- unigram_tokenizer(tok_path)
expect_equal(tok$n_pieces, 32100L)
expect_equal(tok$eos_id, 1L)
expect_equal(tok$pad_id, 0L)
expect_equal(tok$unk_id, 2L)

# --- raw encodings (no padding/truncation) --------------------------------------

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

# --- padding + truncation at max_length -------------------------------------------

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

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

batch <- encode_unigram(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))
