// Copyright 2026 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

// HuggingFace JSON format tests for tokenizer.
//
// These tests load tokenizers from embedded HuggingFace JSON configs and verify
// correct behavior. They complement the builder-API tests in
// tokenizer_encode_test.cc by testing the full JSON loading pipeline with
// realistic configurations.
//
// Buffer/chunk invariance tests verify that tokenizer produces identical
// output across different transform buffer sizes and chunk feeding patterns.

#include <cstring>
#include <string>
#include <vector>

#include "iree/base/api.h"
#include "iree/testing/gtest.h"
#include "iree/testing/status_matchers.h"
#include "iree/tokenizer/format/huggingface/tokenizer_json.h"
#include "iree/tokenizer/testdata/streaming_testdata.h"
#include "iree/tokenizer/tokenizer.h"

namespace iree::tokenizer {
namespace {

//===----------------------------------------------------------------------===//
// Test Utilities
//===----------------------------------------------------------------------===//

static iree_string_view_t GetEmbeddedFile(const char* name) {
  const struct iree_file_toc_t* toc =
      iree_tokenizer_streaming_testdata_create();
  for (size_t i = 0; i < iree_tokenizer_streaming_testdata_size(); ++i) {
    if (strcmp(toc[i].name, name) == 0) {
      return iree_make_string_view(toc[i].data, toc[i].size);
    }
  }
  return iree_string_view_empty();
}

struct TokenizerDeleter {
  void operator()(iree_tokenizer_t* t) const { iree_tokenizer_free(t); }
};
using ScopedTokenizer = std::unique_ptr<iree_tokenizer_t, TokenizerDeleter>;

static iree::StatusOr<ScopedTokenizer> LoadTokenizer(const char* filename) {
  iree_string_view_t json = GetEmbeddedFile(filename);
  if (json.size == 0) {
    return iree_make_status(IREE_STATUS_NOT_FOUND, "file '%s' not in testdata",
                            filename);
  }
  iree_tokenizer_t* tokenizer = nullptr;
  IREE_RETURN_IF_ERROR(iree_tokenizer_from_huggingface_json(
      json, iree_allocator_system(), &tokenizer));
  return ScopedTokenizer(tokenizer);
}

// Encodes text using the one-shot API.
static StatusOr<std::vector<iree_tokenizer_token_id_t>> EncodeOneShot(
    iree_tokenizer_t* tokenizer, iree_string_view_t text) {
  std::vector<iree_tokenizer_token_id_t> ids(256);
  iree_host_size_t count = 0;
  IREE_RETURN_IF_ERROR(
      iree_tokenizer_encode(tokenizer, text, IREE_TOKENIZER_ENCODE_FLAG_NONE,
                            iree_tokenizer_make_token_output(
                                ids.data(), nullptr, nullptr, ids.size()),
                            iree_allocator_system(), &count));
  ids.resize(count);
  return ids;
}

// Encodes text using streaming API with specified transform buffer size.
static StatusOr<std::vector<iree_tokenizer_token_id_t>> EncodeStreaming(
    iree_tokenizer_t* tokenizer, iree_string_view_t text,
    iree_host_size_t buffer_size, iree_host_size_t chunk_size = 0) {
  iree_host_size_t state_size = 0;
  IREE_RETURN_IF_ERROR(
      iree_tokenizer_encode_state_calculate_size(tokenizer, &state_size));

  std::vector<uint8_t> state_storage(state_size);
  std::vector<uint8_t> transform_buffer(buffer_size);

  iree_tokenizer_encode_state_t* state = nullptr;
  IREE_RETURN_IF_ERROR(iree_tokenizer_encode_state_initialize(
      tokenizer,
      iree_make_byte_span(state_storage.data(), state_storage.size()),
      iree_make_byte_span(transform_buffer.data(), transform_buffer.size()),
      iree_tokenizer_offset_run_list_empty(),
      IREE_TOKENIZER_ENCODE_FLAG_AT_INPUT_START, &state));

  std::vector<iree_tokenizer_token_id_t> all_ids;
  std::vector<iree_tokenizer_token_id_t> batch_ids(256);
  iree_host_size_t offset = 0;

  while (offset < text.size) {
    iree_host_size_t remaining = text.size - offset;
    iree_host_size_t feed_size =
        (chunk_size > 0 && chunk_size < remaining) ? chunk_size : remaining;

    iree_string_view_t chunk =
        iree_make_string_view(text.data + offset, feed_size);

    iree_host_size_t consumed = 0;
    iree_host_size_t produced = 0;
    iree_status_t status = iree_tokenizer_encode_state_feed(
        state, chunk,
        iree_tokenizer_make_token_output(batch_ids.data(), nullptr, nullptr,
                                         batch_ids.size()),
        &consumed, &produced);
    if (!iree_status_is_ok(status)) {
      iree_tokenizer_encode_state_deinitialize(state);
      return status;
    }

    for (iree_host_size_t i = 0; i < produced; ++i) {
      all_ids.push_back(batch_ids[i]);
    }
    offset += consumed;

    if (consumed == 0 && produced == 0 && feed_size > 0) {
      iree_tokenizer_encode_state_deinitialize(state);
      return iree_make_status(IREE_STATUS_INTERNAL, "no progress at offset %zu",
                              offset);
    }
  }

  // Drain finalize - may produce multiple batches for large outputs.
  iree_host_size_t finalize_count = 0;
  do {
    iree_status_t status = iree_tokenizer_encode_state_finalize(
        state,
        iree_tokenizer_make_token_output(batch_ids.data(), nullptr, nullptr,
                                         batch_ids.size()),
        &finalize_count);
    if (!iree_status_is_ok(status)) {
      iree_tokenizer_encode_state_deinitialize(state);
      return status;
    }
    for (iree_host_size_t i = 0; i < finalize_count; ++i) {
      all_ids.push_back(batch_ids[i]);
    }
  } while (finalize_count == batch_ids.size());

  iree_tokenizer_encode_state_deinitialize(state);
  return all_ids;
}

//===----------------------------------------------------------------------===//
// BPE ByteLevel Tests
//===----------------------------------------------------------------------===//

class BPEByteLevelTest : public ::testing::Test {
 protected:
  void SetUp() override {
    auto result = LoadTokenizer("bpe_bytelevel_minimal.json");
    if (!result.ok()) {
      GTEST_SKIP() << result.status().ToString();
    }
    tokenizer_ = std::move(result.value());
  }

  iree_tokenizer_t* tokenizer() { return tokenizer_.get(); }

 private:
  ScopedTokenizer tokenizer_;
};

TEST_F(BPEByteLevelTest, LoadsSuccessfully) { ASSERT_NE(tokenizer(), nullptr); }

TEST_F(BPEByteLevelTest, EncodeHelloWorld) {
  IREE_ASSERT_OK_AND_ASSIGN(auto ids,
                            EncodeOneShot(tokenizer(), IREE_SV("hello world")));
  // Correct BPE merge order: "hello" -> ["he", "llo"], " world" -> ["Ġworld"]
  // The merge "l l" at rank 1 fires before "he l" at rank 9, so "hello"
  // tokenizes as ["he", "llo"] not ["hello"].
  EXPECT_EQ(ids.size(), 3u);
  EXPECT_EQ(ids[0], 98);   // he
  EXPECT_EQ(ids[1], 105);  // llo
  EXPECT_EQ(ids[2], 110);  // Ġworld
}

TEST_F(BPEByteLevelTest, BufferInvarianceHelloWorld) {
  iree_string_view_t text = IREE_SV("hello world");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));

  // Test various buffer sizes - all should produce same output.
  for (iree_host_size_t buffer_size : {256, 512, 1024, 4096, 16384}) {
    IREE_ASSERT_OK_AND_ASSIGN(auto result,
                              EncodeStreaming(tokenizer(), text, buffer_size));
    EXPECT_EQ(result, reference)
        << "Buffer size " << buffer_size << " differs from one-shot";
  }
}

TEST_F(BPEByteLevelTest, ChunkInvarianceHelloWorld) {
  iree_string_view_t text = IREE_SV("hello world");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));

  // Test various chunk sizes with fixed buffer.
  for (iree_host_size_t chunk_size : {1, 2, 3, 4, 5, 7, 11}) {
    IREE_ASSERT_OK_AND_ASSIGN(
        auto result, EncodeStreaming(tokenizer(), text, 4096, chunk_size));
    EXPECT_EQ(result, reference)
        << "Chunk size " << chunk_size << " differs from one-shot";
  }
}

TEST_F(BPEByteLevelTest, BufferInvarianceTheQuickBrown) {
  iree_string_view_t text = IREE_SV("the quick brown fox");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));

  for (iree_host_size_t buffer_size : {256, 512, 1024, 4096}) {
    IREE_ASSERT_OK_AND_ASSIGN(auto result,
                              EncodeStreaming(tokenizer(), text, buffer_size));
    EXPECT_EQ(result, reference)
        << "Buffer size " << buffer_size << " differs from one-shot";
  }
}

TEST_F(BPEByteLevelTest, ChunkInvarianceTheQuickBrown) {
  iree_string_view_t text = IREE_SV("the quick brown fox");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));

  for (iree_host_size_t chunk_size : {1, 2, 3, 5, 7}) {
    IREE_ASSERT_OK_AND_ASSIGN(
        auto result, EncodeStreaming(tokenizer(), text, 4096, chunk_size));
    EXPECT_EQ(result, reference)
        << "Chunk size " << chunk_size << " differs from one-shot";
  }
}

TEST_F(BPEByteLevelTest, EmptyInput) {
  IREE_ASSERT_OK_AND_ASSIGN(
      auto ids, EncodeOneShot(tokenizer(), iree_string_view_empty()));
  EXPECT_TRUE(ids.empty());

  IREE_ASSERT_OK_AND_ASSIGN(
      auto streaming,
      EncodeStreaming(tokenizer(), iree_string_view_empty(), 4096));
  EXPECT_TRUE(streaming.empty());
}

TEST_F(BPEByteLevelTest, SingleCharacter) {
  iree_string_view_t text = IREE_SV("a");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));

  IREE_ASSERT_OK_AND_ASSIGN(auto streaming,
                            EncodeStreaming(tokenizer(), text, 256));
  EXPECT_EQ(streaming, reference);
}

TEST_F(BPEByteLevelTest, RepeatedCharacters) {
  iree_string_view_t text = IREE_SV("aaaaaaaaaa");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));

  for (iree_host_size_t buffer_size : {256, 512, 1024}) {
    IREE_ASSERT_OK_AND_ASSIGN(auto result,
                              EncodeStreaming(tokenizer(), text, buffer_size));
    EXPECT_EQ(result, reference)
        << "Buffer size " << buffer_size << " differs for repeated chars";
  }
}

TEST_F(BPEByteLevelTest, MultipleSpaces) {
  iree_string_view_t text = IREE_SV("hello  world");  // Two spaces.
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));

  for (iree_host_size_t chunk_size : {1, 2, 3}) {
    IREE_ASSERT_OK_AND_ASSIGN(
        auto result, EncodeStreaming(tokenizer(), text, 4096, chunk_size));
    EXPECT_EQ(result, reference)
        << "Chunk size " << chunk_size << " differs for multiple spaces";
  }
}

// Tests that " the" correctly merges through the chain: Ġ+t -> Ġt -> Ġth ->
// Ġthe even when the space and word are split across chunk boundaries.
TEST_F(BPEByteLevelTest, SpacePrefixedWordMerge) {
  iree_string_view_t text = IREE_SV("hello the world");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));

  // Chunk sizes that force " the" to split at different points:
  // size 5: "hello" | " the " | "world" - space at chunk boundary
  // size 6: "hello " | "the wo" | "rld" - after space
  // size 7: "hello t" | "he worl" | "d" - mid-word
  for (iree_host_size_t chunk_size : {1, 2, 3, 5, 6, 7}) {
    IREE_ASSERT_OK_AND_ASSIGN(
        auto result, EncodeStreaming(tokenizer(), text, 4096, chunk_size));
    EXPECT_EQ(result, reference)
        << "Chunk size " << chunk_size << " differs for space-prefixed word";
  }
}

// Tests special token recognition when the token spans chunk boundaries.
// "<|endoftext|>" is 13 characters - splitting it should still produce one
// token.
TEST_F(BPEByteLevelTest, SpecialTokenSpanning) {
  iree_string_view_t text = IREE_SV("<|endoftext|>");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));
  // Special token should be recognized as a single token (ID 111).
  EXPECT_EQ(reference.size(), 1u);
  EXPECT_EQ(reference[0], 111);

  // Feed the special token character by character.
  for (iree_host_size_t chunk_size : {1, 2, 3, 4, 5, 6, 7}) {
    IREE_ASSERT_OK_AND_ASSIGN(
        auto result, EncodeStreaming(tokenizer(), text, 4096, chunk_size));
    EXPECT_EQ(result, reference)
        << "Chunk size " << chunk_size << " differs for special token spanning";
  }
}

//===----------------------------------------------------------------------===//
// Special Token Negative Tests
//===----------------------------------------------------------------------===//

// Tests that text around special tokens doesn't incorrectly match.
// "x<|endoftext|>y" should NOT match the special token.
TEST_F(BPEByteLevelTest, SpecialTokenNegative) {
  iree_string_view_t text = IREE_SV("x<|endoftext|>y");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));
  // Should NOT be a single token - the surrounding chars prevent special match.
  EXPECT_GT(reference.size(), 1u);

  for (iree_host_size_t chunk_size : {1, 2, 3}) {
    IREE_ASSERT_OK_AND_ASSIGN(
        auto result, EncodeStreaming(tokenizer(), text, 4096, chunk_size));
    EXPECT_EQ(result, reference)
        << "Chunk size " << chunk_size << " differs for special token negative";
  }
}

// Tests behavior at exact buffer boundaries - fill buffer exactly then
// overflow.
TEST_F(BPEByteLevelTest, BufferBoundaryExact) {
  // 64 'a' characters to test exact buffer fill and slight overflow.
  std::string text_64(64, 'a');
  std::string text_65(65, 'a');
  std::string text_128(128, 'a');
  std::string text_129(129, 'a');

  for (const auto& text : {text_64, text_65, text_128, text_129}) {
    iree_string_view_t sv = iree_make_string_view(text.data(), text.size());
    IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), sv));

    // Test with buffer sizes that match or nearly match input size.
    for (iree_host_size_t buffer_size : {64, 128, 256}) {
      IREE_ASSERT_OK_AND_ASSIGN(auto result,
                                EncodeStreaming(tokenizer(), sv, buffer_size));
      EXPECT_EQ(result, reference)
          << "Buffer size " << buffer_size << " differs for " << text.size()
          << " chars";
    }
  }
}

// Tests streaming with empty chunks interspersed between content.
TEST_F(BPEByteLevelTest, EmptyChunks) {
  iree_string_view_t text = IREE_SV("hello");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));

  // Feed character by character - the implementation should handle this.
  IREE_ASSERT_OK_AND_ASSIGN(auto streaming,
                            EncodeStreaming(tokenizer(), text, 4096, 1));
  EXPECT_EQ(streaming, reference);
}

//===----------------------------------------------------------------------===//
// Post-Normalization Special Token Tests
//===----------------------------------------------------------------------===//

// Post-normalization special tokens have normalized=true and are matched AFTER
// the normalizer runs. This test suite verifies correct behavior when these
// tokens appear in various positions and streaming scenarios.

class PostNormSpecialTokenTest : public ::testing::Test {
 protected:
  void SetUp() override {
    // Minimal BPE tokenizer with:
    // - Regular vocab tokens for basic text
    // - A post-normalization special token (normalized=true)
    const char* json = R"({
      "model": {
        "type": "BPE",
        "vocab": {
          "h": 0, "e": 1, "l": 2, "o": 3, " ": 4, "w": 5, "r": 6, "d": 7,
          "he": 8, "ll": 9, "lo": 10,
          "<|special|>": 100
        },
        "merges": ["h e", "l l", "l o"]
      },
      "added_tokens": [{
        "id": 100,
        "content": "<|special|>",
        "single_word": false,
        "lstrip": false,
        "rstrip": false,
        "normalized": true,
        "special": true
      }],
      "pre_tokenizer": {"type": "ByteLevel", "add_prefix_space": false, "trim_offsets": false, "use_regex": true}
    })";
    iree_tokenizer_t* tokenizer = nullptr;
    iree_status_t status = iree_tokenizer_from_huggingface_json(
        iree_make_string_view(json, strlen(json)), iree_allocator_system(),
        &tokenizer);
    IREE_CHECK_OK(status);
    tokenizer_.reset(tokenizer);
  }

  iree_tokenizer_t* tokenizer() { return tokenizer_.get(); }

 private:
  ScopedTokenizer tokenizer_;
};

// Verifies that post-norm special tokens appearing after regular text are
// correctly matched. When streaming, the special token may be encountered while
// segments from earlier text are still pending for the model. The
// implementation must not corrupt ring buffer state by advancing read_position
// prematurely.
TEST_F(PostNormSpecialTokenTest, SpecialTokenAfterText) {
  iree_string_view_t text = IREE_SV("hello <|special|>");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));

  // The last token should be the special token (id=100).
  EXPECT_EQ(reference.back(), 100u);

  // Streaming with various chunk sizes must produce identical results.
  for (iree_host_size_t chunk_size : {1, 2, 3, 5, 7}) {
    IREE_ASSERT_OK_AND_ASSIGN(
        auto result, EncodeStreaming(tokenizer(), text, 4096, chunk_size));
    EXPECT_EQ(result, reference)
        << "Chunk size " << chunk_size
        << " produced different result for post-norm special token after text";
  }
}

// Verifies that post-norm special tokens at the start of input work correctly.
TEST_F(PostNormSpecialTokenTest, SpecialTokenAtStart) {
  iree_string_view_t text = IREE_SV("<|special|>hello");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));

  // The first token should be the special token (id=100).
  EXPECT_EQ(reference.front(), 100u);

  for (iree_host_size_t chunk_size : {1, 2, 3}) {
    IREE_ASSERT_OK_AND_ASSIGN(
        auto result, EncodeStreaming(tokenizer(), text, 4096, chunk_size));
    EXPECT_EQ(result, reference)
        << "Chunk size " << chunk_size
        << " produced different result for post-norm special token at start";
  }
}

// Verifies that multiple post-norm special tokens in sequence work correctly.
TEST_F(PostNormSpecialTokenTest, MultipleSpecialTokens) {
  iree_string_view_t text = IREE_SV("<|special|><|special|>");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));

  // Should be exactly two special tokens.
  EXPECT_EQ(reference.size(), 2u);
  EXPECT_EQ(reference[0], 100u);
  EXPECT_EQ(reference[1], 100u);

  for (iree_host_size_t chunk_size : {1, 2, 3}) {
    IREE_ASSERT_OK_AND_ASSIGN(
        auto result, EncodeStreaming(tokenizer(), text, 4096, chunk_size));
    EXPECT_EQ(result, reference)
        << "Chunk size " << chunk_size
        << " produced different result for multiple post-norm special tokens";
  }
}

// Verifies that streaming with very small buffers handles post-norm tokens
// correctly when interleaved with regular text. This stresses the ring buffer
// management when segments are pending.
TEST_F(PostNormSpecialTokenTest, SmallBufferWithInterleavedTokens) {
  iree_string_view_t text = IREE_SV("hello <|special|> world");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));

  // Use a small buffer to stress the ring buffer management.
  // The special token should appear somewhere in the middle.
  bool found_special = false;
  for (auto id : reference) {
    if (id == 100u) found_special = true;
  }
  EXPECT_TRUE(found_special) << "Special token not found in reference output";

  // Small buffer sizes stress the segment queuing logic.
  for (iree_host_size_t buffer_size : {64, 128, 256}) {
    IREE_ASSERT_OK_AND_ASSIGN(
        auto result, EncodeStreaming(tokenizer(), text, buffer_size, 1));
    EXPECT_EQ(result, reference)
        << "Buffer size " << buffer_size
        << " with 1-byte chunks produced different result";
  }
}

//===----------------------------------------------------------------------===//
// Whitespace Pre-Tokenizer Tests
//===----------------------------------------------------------------------===//

// Tests that the Whitespace pre_tokenizer works correctly with streaming
// encoding, including proper word boundary detection across chunk boundaries.

class WhitespacePreTokenizerTest : public ::testing::Test {
 protected:
  void SetUp() override {
    // BPE tokenizer with Whitespace pre_tokenizer.
    // Vocab and merges must be consistent: every merge operand must exist.
    // "hello" merges: h+e→he, l+l→ll, ll+o→llo, he+llo→hello
    // "world" merges: w+o→wo, r+l→rl, rl+d→rld, wo+rld→world
    const char* json = R"({
      "model": {
        "type": "BPE",
        "vocab": {
          "h": 0, "e": 1, "l": 2, "o": 3, "w": 4, "r": 5, "d": 6,
          "he": 10, "ll": 11, "wo": 12, "rl": 13,
          "llo": 20, "rld": 21,
          "hello": 30, "world": 31
        },
        "merges": [
          "h e", "l l", "w o", "r l",
          "ll o", "rl d",
          "he llo", "wo rld"
        ]
      },
      "pre_tokenizer": {"type": "Whitespace"}
    })";
    iree_tokenizer_t* tokenizer = nullptr;
    iree_status_t status = iree_tokenizer_from_huggingface_json(
        iree_make_string_view(json, strlen(json)), iree_allocator_system(),
        &tokenizer);
    IREE_CHECK_OK(status);
    tokenizer_.reset(tokenizer);
  }

  iree_tokenizer_t* tokenizer() { return tokenizer_.get(); }

 private:
  ScopedTokenizer tokenizer_;
};

TEST_F(WhitespacePreTokenizerTest, LoadsSuccessfully) {
  EXPECT_NE(tokenizer(), nullptr);
}

TEST_F(WhitespacePreTokenizerTest, SingleWord) {
  iree_string_view_t text = IREE_SV("hello");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));
  // "hello" should merge to a single token (ID 30).
  EXPECT_EQ(reference.size(), 1u);
  EXPECT_EQ(reference[0], 30);
}

TEST_F(WhitespacePreTokenizerTest, TwoWords) {
  iree_string_view_t text = IREE_SV("hello world");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));
  // "hello world" -> ["hello", "world"] -> [30, 31].
  EXPECT_EQ(reference.size(), 2u);
  EXPECT_EQ(reference[0], 30);
  EXPECT_EQ(reference[1], 31);
}

TEST_F(WhitespacePreTokenizerTest, StreamingChunkInvariance) {
  iree_string_view_t text = IREE_SV("hello world");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));

  // Test various chunk sizes that split at different positions:
  // "h|ello world", "he|llo world", "hel|lo world", "hell|o world",
  // "hello| world", "hello |world", etc.
  for (iree_host_size_t chunk_size = 1; chunk_size <= 11; ++chunk_size) {
    IREE_ASSERT_OK_AND_ASSIGN(
        auto result, EncodeStreaming(tokenizer(), text, 4096, chunk_size));
    EXPECT_EQ(result, reference)
        << "Chunk size " << chunk_size << " produced different result";
  }
}

TEST_F(WhitespacePreTokenizerTest, StreamingBufferInvariance) {
  iree_string_view_t text = IREE_SV("hello world");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));

  // Test various buffer sizes that affect ring buffer behavior.
  for (iree_host_size_t buffer_size : {16, 32, 64, 128, 256, 4096}) {
    IREE_ASSERT_OK_AND_ASSIGN(auto result,
                              EncodeStreaming(tokenizer(), text, buffer_size));
    EXPECT_EQ(result, reference)
        << "Buffer size " << buffer_size << " produced different result";
  }
}

TEST_F(WhitespacePreTokenizerTest, MultipleSpaces) {
  iree_string_view_t text = IREE_SV("hello    world");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));
  // Multiple spaces between words should still produce two tokens.
  EXPECT_EQ(reference.size(), 2u);

  for (iree_host_size_t chunk_size : {1, 2, 3, 5}) {
    IREE_ASSERT_OK_AND_ASSIGN(
        auto result, EncodeStreaming(tokenizer(), text, 4096, chunk_size));
    EXPECT_EQ(result, reference)
        << "Chunk size " << chunk_size << " differs for multiple spaces";
  }
}

TEST_F(WhitespacePreTokenizerTest, LeadingAndTrailingWhitespace) {
  iree_string_view_t text = IREE_SV("  hello world  ");
  IREE_ASSERT_OK_AND_ASSIGN(auto reference, EncodeOneShot(tokenizer(), text));
  // Leading and trailing whitespace should be ignored.
  EXPECT_EQ(reference.size(), 2u);

  for (iree_host_size_t chunk_size : {1, 2, 3, 5}) {
    IREE_ASSERT_OK_AND_ASSIGN(
        auto result, EncodeStreaming(tokenizer(), text, 4096, chunk_size));
    EXPECT_EQ(result, reference) << "Chunk size " << chunk_size
                                 << " differs for leading/trailing whitespace";
  }
}

TEST_F(WhitespacePreTokenizerTest, WhitespaceSplitAlias) {
  // WhitespaceSplit is an alias for Whitespace - verify it loads successfully.
  const char* json = R"({
    "model": {
      "type": "BPE",
      "vocab": {"a": 0},
      "merges": []
    },
    "pre_tokenizer": {"type": "WhitespaceSplit"}
  })";
  iree_tokenizer_t* tokenizer = nullptr;
  iree_status_t status = iree_tokenizer_from_huggingface_json(
      iree_make_string_view(json, strlen(json)), iree_allocator_system(),
      &tokenizer);
  IREE_EXPECT_OK(status);
  if (tokenizer) iree_tokenizer_free(tokenizer);
}

//===----------------------------------------------------------------------===//
// HuggingFace Ground Truth Validation Tests
//===----------------------------------------------------------------------===//
// These tests verify that IREE's tokenizer produces IDENTICAL output to
// HuggingFace across multiple buffer sizes. Expected token IDs are generated
// from the HuggingFace tokenizers library:
//   uv run --with tokenizers python generate_testdata_expected_ids.py

// Buffer sizes to test. Smaller sizes force partial-segment mode.
constexpr iree_host_size_t kGroundTruthBufferSizes[] = {64, 128, 256, 512,
                                                        65536};

// Expected token IDs from HuggingFace tokenizers library.
// "hello" -> ['he', 'llo']
static constexpr iree_tokenizer_token_id_t kHelloExpected[] = {98, 105};

// "hello world" -> ['he', 'llo', 'Ġworld']
static constexpr iree_tokenizer_token_id_t kHelloWorldExpected[] = {98, 105,
                                                                    110};

// "The quick brown fox" -> individual chars with Ġ for spaces
static constexpr iree_tokenizer_token_id_t kSimpleASCIIExpected[] = {
    51, 98, 94, 80, 84, 72, 66, 74, 94, 65, 81, 78, 86, 77, 94, 69, 78, 87};

// "hello<|endoftext|>world" -> ['he', 'llo', '<|endoftext|>', 'w', 'orld']
static constexpr iree_tokenizer_token_id_t kWithSpecialExpected[] = {
    98, 105, 111, 86, 108};

// "  hello   world  " -> spaces become Ġ tokens
static constexpr iree_tokenizer_token_id_t kWhitespaceExpected[] = {
    94, 94, 98, 105, 94, 94, 110, 94, 94};

// "a" -> ['a']
static constexpr iree_tokenizer_token_id_t kSingleCharExpected[] = {64};

// "12345" -> ['1', '2', '3', '4', '5']
static constexpr iree_tokenizer_token_id_t kNumbersExpected[] = {16, 17, 18, 19,
                                                                 20};

// "!@#$%^&*()" -> individual punctuation tokens
static constexpr iree_tokenizer_token_id_t kPunctuationExpected[] = {
    0, 31, 2, 3, 4, 61, 5, 9, 7, 8};

// Parameterized test fixture for ground truth validation.
class GroundTruthTest : public ::testing::TestWithParam<iree_host_size_t> {
 protected:
  void SetUp() override {
    iree_string_view_t json_str = GetEmbeddedFile("bpe_bytelevel_minimal.json");
    ASSERT_GT(json_str.size, 0u) << "Tokenizer JSON not found";
    IREE_ASSERT_OK(iree_tokenizer_from_huggingface_json(
        json_str, iree_allocator_system(), &tokenizer_));
  }

  void TearDown() override {
    if (tokenizer_) {
      iree_tokenizer_free(tokenizer_);
    }
  }

  // Encode input text using the current buffer size (from parameterized test).
  std::vector<iree_tokenizer_token_id_t> EncodeWithBufferSize(
      const char* input, iree_host_size_t buffer_size) {
    iree_host_size_t state_size = 0;
    IREE_EXPECT_OK(
        iree_tokenizer_encode_state_calculate_size(tokenizer_, &state_size));

    std::vector<uint8_t> state_storage(state_size);
    std::vector<uint8_t> transform_buffer(buffer_size);

    iree_tokenizer_encode_state_t* state = nullptr;
    IREE_EXPECT_OK(iree_tokenizer_encode_state_initialize(
        tokenizer_,
        iree_make_byte_span(state_storage.data(), state_storage.size()),
        iree_make_byte_span(transform_buffer.data(), transform_buffer.size()),
        iree_tokenizer_offset_run_list_empty(),
        IREE_TOKENIZER_ENCODE_FLAG_AT_INPUT_START, &state));

    std::vector<iree_tokenizer_token_id_t> tokens(1024);
    iree_host_size_t total_tokens = 0;

    iree_string_view_t input_view = iree_make_string_view(input, strlen(input));
    iree_host_size_t bytes_consumed = 0;
    iree_host_size_t token_count = 0;

    IREE_EXPECT_OK(iree_tokenizer_encode_state_feed(
        state, input_view,
        iree_tokenizer_make_token_output(tokens.data(), NULL, NULL,
                                         tokens.size()),
        &bytes_consumed, &token_count));
    total_tokens += token_count;

    iree_host_size_t finalize_count = 0;
    IREE_EXPECT_OK(iree_tokenizer_encode_state_finalize(
        state,
        iree_tokenizer_make_token_output(tokens.data() + total_tokens, NULL,
                                         NULL, tokens.size() - total_tokens),
        &finalize_count));
    total_tokens += finalize_count;

    tokens.resize(total_tokens);
    iree_tokenizer_encode_state_deinitialize(state);

    return tokens;
  }

  // Encode using byte-by-byte feeding (most stressful streaming test).
  std::vector<iree_tokenizer_token_id_t> EncodeByteByByte(
      const char* input, iree_host_size_t buffer_size) {
    iree_host_size_t state_size = 0;
    IREE_EXPECT_OK(
        iree_tokenizer_encode_state_calculate_size(tokenizer_, &state_size));

    std::vector<uint8_t> state_storage(state_size);
    std::vector<uint8_t> transform_buffer(buffer_size);

    iree_tokenizer_encode_state_t* state = nullptr;
    IREE_EXPECT_OK(iree_tokenizer_encode_state_initialize(
        tokenizer_,
        iree_make_byte_span(state_storage.data(), state_storage.size()),
        iree_make_byte_span(transform_buffer.data(), transform_buffer.size()),
        iree_tokenizer_offset_run_list_empty(),
        IREE_TOKENIZER_ENCODE_FLAG_AT_INPUT_START, &state));

    std::vector<iree_tokenizer_token_id_t> tokens(1024);
    iree_host_size_t total_tokens = 0;
    size_t input_length = strlen(input);

    for (size_t i = 0; i < input_length; ++i) {
      iree_string_view_t chunk = iree_make_string_view(input + i, 1);
      iree_host_size_t bytes_consumed = 0;
      iree_host_size_t token_count = 0;

      IREE_EXPECT_OK(iree_tokenizer_encode_state_feed(
          state, chunk,
          iree_tokenizer_make_token_output(tokens.data() + total_tokens, NULL,
                                           NULL, tokens.size() - total_tokens),
          &bytes_consumed, &token_count));
      total_tokens += token_count;
    }

    iree_host_size_t finalize_count = 0;
    IREE_EXPECT_OK(iree_tokenizer_encode_state_finalize(
        state,
        iree_tokenizer_make_token_output(tokens.data() + total_tokens, NULL,
                                         NULL, tokens.size() - total_tokens),
        &finalize_count));
    total_tokens += finalize_count;

    tokens.resize(total_tokens);
    iree_tokenizer_encode_state_deinitialize(state);

    return tokens;
  }

  iree_tokenizer_t* tokenizer_ = nullptr;
};

TEST_P(GroundTruthTest, HelloWorld) {
  iree_host_size_t buffer_size = GetParam();
  auto tokens = EncodeWithBufferSize("hello world", buffer_size);

  std::vector<iree_tokenizer_token_id_t> expected(
      kHelloWorldExpected,
      kHelloWorldExpected +
          sizeof(kHelloWorldExpected) / sizeof(kHelloWorldExpected[0]));

  EXPECT_EQ(tokens, expected)
      << "Buffer size: " << buffer_size << ", input: 'hello world'";
}

TEST_P(GroundTruthTest, HelloWorldByteByByte) {
  iree_host_size_t buffer_size = GetParam();
  auto tokens = EncodeByteByByte("hello world", buffer_size);

  std::vector<iree_tokenizer_token_id_t> expected(
      kHelloWorldExpected,
      kHelloWorldExpected +
          sizeof(kHelloWorldExpected) / sizeof(kHelloWorldExpected[0]));

  EXPECT_EQ(tokens, expected) << "Buffer size: " << buffer_size
                              << ", input: 'hello world' (byte-by-byte)";
}

TEST_P(GroundTruthTest, SimpleASCII) {
  iree_host_size_t buffer_size = GetParam();
  auto tokens = EncodeWithBufferSize("The quick brown fox", buffer_size);

  std::vector<iree_tokenizer_token_id_t> expected(
      kSimpleASCIIExpected,
      kSimpleASCIIExpected +
          sizeof(kSimpleASCIIExpected) / sizeof(kSimpleASCIIExpected[0]));

  EXPECT_EQ(tokens, expected)
      << "Buffer size: " << buffer_size << ", input: 'The quick brown fox'";
}

TEST_P(GroundTruthTest, WithSpecialToken) {
  iree_host_size_t buffer_size = GetParam();
  auto tokens = EncodeWithBufferSize("hello<|endoftext|>world", buffer_size);

  std::vector<iree_tokenizer_token_id_t> expected(
      kWithSpecialExpected,
      kWithSpecialExpected +
          sizeof(kWithSpecialExpected) / sizeof(kWithSpecialExpected[0]));

  EXPECT_EQ(tokens, expected)
      << "Buffer size: " << buffer_size << ", input: 'hello<|endoftext|>world'";
}

TEST_P(GroundTruthTest, Whitespace) {
  iree_host_size_t buffer_size = GetParam();
  auto tokens = EncodeWithBufferSize("  hello   world  ", buffer_size);

  std::vector<iree_tokenizer_token_id_t> expected(
      kWhitespaceExpected,
      kWhitespaceExpected +
          sizeof(kWhitespaceExpected) / sizeof(kWhitespaceExpected[0]));

  EXPECT_EQ(tokens, expected)
      << "Buffer size: " << buffer_size << ", input: '  hello   world  '";
}

TEST_P(GroundTruthTest, SingleChar) {
  iree_host_size_t buffer_size = GetParam();
  auto tokens = EncodeWithBufferSize("a", buffer_size);

  std::vector<iree_tokenizer_token_id_t> expected(
      kSingleCharExpected,
      kSingleCharExpected +
          sizeof(kSingleCharExpected) / sizeof(kSingleCharExpected[0]));

  EXPECT_EQ(tokens, expected)
      << "Buffer size: " << buffer_size << ", input: 'a'";
}

TEST_P(GroundTruthTest, Numbers) {
  iree_host_size_t buffer_size = GetParam();
  auto tokens = EncodeWithBufferSize("12345", buffer_size);

  std::vector<iree_tokenizer_token_id_t> expected(
      kNumbersExpected, kNumbersExpected + sizeof(kNumbersExpected) /
                                               sizeof(kNumbersExpected[0]));

  EXPECT_EQ(tokens, expected)
      << "Buffer size: " << buffer_size << ", input: '12345'";
}

TEST_P(GroundTruthTest, Punctuation) {
  iree_host_size_t buffer_size = GetParam();
  auto tokens = EncodeWithBufferSize("!@#$%^&*()", buffer_size);

  std::vector<iree_tokenizer_token_id_t> expected(
      kPunctuationExpected,
      kPunctuationExpected +
          sizeof(kPunctuationExpected) / sizeof(kPunctuationExpected[0]));

  EXPECT_EQ(tokens, expected)
      << "Buffer size: " << buffer_size << ", input: '!@#$%^&*()'";
}

TEST_P(GroundTruthTest, EmptyInput) {
  iree_host_size_t buffer_size = GetParam();
  auto tokens = EncodeWithBufferSize("", buffer_size);

  EXPECT_TRUE(tokens.empty()) << "Buffer size: " << buffer_size
                              << ", expected empty output for empty input";
}

// Run all parameterized tests with each buffer size.
INSTANTIATE_TEST_SUITE_P(
    BufferSizes, GroundTruthTest, ::testing::ValuesIn(kGroundTruthBufferSizes),
    [](const ::testing::TestParamInfo<iree_host_size_t>& info) {
      return "BufferSize" + std::to_string(info.param);
    });

// Helper to format strings for test output.
static std::string repr(const char* s) {
  std::string result = "'";
  for (const char* p = s; *p; ++p) {
    if (*p == '\n')
      result += "\\n";
    else if (*p == '\t')
      result += "\\t";
    else if (*p == '\'')
      result += "\\'";
    else if (*p < 32)
      result += "\\x" + std::to_string((unsigned char)*p);
    else
      result += *p;
  }
  result += "'";
  return result;
}

// Verifies that ALL buffer sizes produce IDENTICAL output.
TEST(GroundTruthInvarianceTest, AllBufferSizesProduceIdenticalOutput) {
  iree_string_view_t json_str = GetEmbeddedFile("bpe_bytelevel_minimal.json");
  ASSERT_GT(json_str.size, 0u);

  iree_tokenizer_t* tokenizer = nullptr;
  IREE_ASSERT_OK(iree_tokenizer_from_huggingface_json(
      json_str, iree_allocator_system(), &tokenizer));

  const char* test_inputs[] = {
      "hello world",
      "The quick brown fox jumps over the lazy dog",
      "hello<|endoftext|>world",
      "  lots   of   whitespace  ",
      "!@#$%^&*()",
      "12345",
      "",
  };

  for (const char* input : test_inputs) {
    std::vector<std::vector<iree_tokenizer_token_id_t>> results_per_size;

    for (iree_host_size_t buffer_size : kGroundTruthBufferSizes) {
      iree_host_size_t state_size = 0;
      IREE_EXPECT_OK(
          iree_tokenizer_encode_state_calculate_size(tokenizer, &state_size));

      std::vector<uint8_t> state_storage(state_size);
      std::vector<uint8_t> transform_buffer(buffer_size);

      iree_tokenizer_encode_state_t* state = nullptr;
      IREE_EXPECT_OK(iree_tokenizer_encode_state_initialize(
          tokenizer,
          iree_make_byte_span(state_storage.data(), state_storage.size()),
          iree_make_byte_span(transform_buffer.data(), transform_buffer.size()),
          iree_tokenizer_offset_run_list_empty(),
          IREE_TOKENIZER_ENCODE_FLAG_AT_INPUT_START, &state));

      std::vector<iree_tokenizer_token_id_t> tokens(1024);
      iree_host_size_t total_tokens = 0;

      iree_string_view_t input_view =
          iree_make_string_view(input, strlen(input));
      iree_host_size_t bytes_consumed = 0;
      iree_host_size_t token_count = 0;

      IREE_EXPECT_OK(iree_tokenizer_encode_state_feed(
          state, input_view,
          iree_tokenizer_make_token_output(tokens.data(), NULL, NULL,
                                           tokens.size()),
          &bytes_consumed, &token_count));
      total_tokens += token_count;

      iree_host_size_t finalize_count = 0;
      IREE_EXPECT_OK(iree_tokenizer_encode_state_finalize(
          state,
          iree_tokenizer_make_token_output(tokens.data() + total_tokens, NULL,
                                           NULL, tokens.size() - total_tokens),
          &finalize_count));
      total_tokens += finalize_count;

      tokens.resize(total_tokens);
      iree_tokenizer_encode_state_deinitialize(state);

      results_per_size.push_back(tokens);
    }

    // Verify all buffer sizes produced identical results.
    for (size_t i = 1; i < results_per_size.size(); ++i) {
      EXPECT_EQ(results_per_size[0], results_per_size[i])
          << "Input: " << repr(input) << "\n"
          << "Buffer size " << kGroundTruthBufferSizes[0] << " produced "
          << results_per_size[0].size() << " tokens\n"
          << "Buffer size " << kGroundTruthBufferSizes[i] << " produced "
          << results_per_size[i].size() << " tokens";
    }
  }

  iree_tokenizer_free(tokenizer);
}

//===----------------------------------------------------------------------===//
// Pair Encoding Tests (BERT TemplateProcessing)
//===----------------------------------------------------------------------===//

// Helper: batch-encode a single (text, text_pair) item with ADD_SPECIAL_TOKENS.
// Returns {token_ids, type_ids} vectors.
struct PairEncodeResult {
  std::vector<iree_tokenizer_token_id_t> token_ids;
  std::vector<uint8_t> type_ids;
};

struct PairEncodeDetailedResult {
  std::vector<iree_tokenizer_token_id_t> token_ids;
  std::vector<uint8_t> type_ids;
  std::vector<iree_tokenizer_offset_t> token_offsets;
};

static StatusOr<PairEncodeResult> EncodePair(iree_tokenizer_t* tokenizer,
                                             iree_string_view_t text,
                                             iree_string_view_t text_pair) {
  iree_host_size_t state_size = 0;
  IREE_RETURN_IF_ERROR(
      iree_tokenizer_encode_state_calculate_size(tokenizer, &state_size));
  std::vector<uint8_t> state_storage(state_size);
  std::vector<uint8_t> transform_buffer(4096);

  std::vector<iree_tokenizer_token_id_t> token_ids(256);
  std::vector<uint8_t> type_ids(256, 0xFF);

  iree_tokenizer_encode_batch_item_t item = {};
  item.text = text;
  item.text_pair = text_pair;
  item.flags = IREE_TOKENIZER_ENCODE_BATCH_ITEM_FLAG_HAS_TEXT_PAIR;
  item.output = iree_tokenizer_make_token_output(token_ids.data(), nullptr,
                                                 type_ids.data(), 256);
  item.out_token_count = 0;

  IREE_RETURN_IF_ERROR(iree_tokenizer_encode_batch(
      tokenizer, &item, 1, IREE_TOKENIZER_ENCODE_FLAG_ADD_SPECIAL_TOKENS,
      iree_make_byte_span(state_storage.data(), state_storage.size()),
      iree_make_byte_span(transform_buffer.data(), transform_buffer.size()),
      iree_tokenizer_offset_run_list_empty()));

  token_ids.resize(item.out_token_count);
  type_ids.resize(item.out_token_count);
  return PairEncodeResult{std::move(token_ids), std::move(type_ids)};
}

static StatusOr<PairEncodeDetailedResult> EncodePairDetailed(
    iree_tokenizer_t* tokenizer, iree_string_view_t text,
    iree_string_view_t text_pair, iree_host_size_t output_capacity,
    bool track_offsets) {
  iree_host_size_t state_size = 0;
  IREE_RETURN_IF_ERROR(
      iree_tokenizer_encode_state_calculate_size(tokenizer, &state_size));
  std::vector<uint8_t> state_storage(state_size);
  std::vector<uint8_t> transform_buffer(4096);
  std::vector<iree_tokenizer_offset_run_t> offset_runs(256);

  iree_host_size_t storage_capacity = output_capacity > 0 ? output_capacity : 1;
  std::vector<iree_tokenizer_token_id_t> token_ids(storage_capacity);
  std::vector<uint8_t> type_ids(storage_capacity, 0xFF);
  std::vector<iree_tokenizer_offset_t> token_offsets(storage_capacity);

  iree_tokenizer_encode_batch_item_t item = {};
  item.text = text;
  item.text_pair = text_pair;
  item.flags = IREE_TOKENIZER_ENCODE_BATCH_ITEM_FLAG_HAS_TEXT_PAIR;
  item.output = iree_tokenizer_make_token_output(
      token_ids.data(), track_offsets ? token_offsets.data() : nullptr,
      type_ids.data(), output_capacity);

  iree_tokenizer_offset_run_list_t offset_run_list = {
      /*.capacity=*/offset_runs.size(),
      /*.values=*/offset_runs.data(),
  };
  IREE_RETURN_IF_ERROR(iree_tokenizer_encode_batch(
      tokenizer, &item, 1, IREE_TOKENIZER_ENCODE_FLAG_ADD_SPECIAL_TOKENS,
      iree_make_byte_span(state_storage.data(), state_storage.size()),
      iree_make_byte_span(transform_buffer.data(), transform_buffer.size()),
      track_offsets ? offset_run_list
                    : iree_tokenizer_offset_run_list_empty()));

  token_ids.resize(item.out_token_count);
  type_ids.resize(item.out_token_count);
  token_offsets.resize(track_offsets ? item.out_token_count : 0);
  return PairEncodeDetailedResult{std::move(token_ids), std::move(type_ids),
                                  std::move(token_offsets)};
}

// Helper: batch-encode single-sequence with ADD_SPECIAL_TOKENS.
static StatusOr<std::vector<iree_tokenizer_token_id_t>> EncodeSingle(
    iree_tokenizer_t* tokenizer, iree_string_view_t text) {
  iree_host_size_t state_size = 0;
  IREE_RETURN_IF_ERROR(
      iree_tokenizer_encode_state_calculate_size(tokenizer, &state_size));
  std::vector<uint8_t> state_storage(state_size);
  std::vector<uint8_t> transform_buffer(4096);

  std::vector<iree_tokenizer_token_id_t> token_ids(256);

  iree_tokenizer_encode_batch_item_t item = {};
  item.text = text;
  item.output =
      iree_tokenizer_make_token_output(token_ids.data(), nullptr, nullptr, 256);
  item.out_token_count = 0;

  IREE_RETURN_IF_ERROR(iree_tokenizer_encode_batch(
      tokenizer, &item, 1, IREE_TOKENIZER_ENCODE_FLAG_ADD_SPECIAL_TOKENS,
      iree_make_byte_span(state_storage.data(), state_storage.size()),
      iree_make_byte_span(transform_buffer.data(), transform_buffer.size()),
      iree_tokenizer_offset_run_list_empty()));

  token_ids.resize(item.out_token_count);
  return token_ids;
}

TEST(PairEncode, BertSingleSequence) {
  // Verify single-sequence encoding still works with the BERT pair tokenizer.
  // Expected: [CLS]=101, "hello"=7592, [SEP]=102
  IREE_ASSERT_OK_AND_ASSIGN(auto tokenizer,
                            LoadTokenizer("wordpiece_bert_pair_minimal.json"));

  IREE_ASSERT_OK_AND_ASSIGN(auto ids,
                            EncodeSingle(tokenizer.get(), IREE_SV("hello")));

  ASSERT_EQ(ids.size(), 3u);
  EXPECT_EQ(ids[0], 101);   // [CLS]
  EXPECT_EQ(ids[1], 7592);  // hello
  EXPECT_EQ(ids[2], 102);   // [SEP]
}

TEST(PairEncode, BertPairSequence) {
  // BERT pair: [CLS] query [SEP] document [SEP]
  // Input: text="hello", text_pair="world"
  // Expected token_ids: [101, 7592, 102, 2088, 102]
  // Expected type_ids:  [  0,    0,   0,    1,   1]
  IREE_ASSERT_OK_AND_ASSIGN(auto tokenizer,
                            LoadTokenizer("wordpiece_bert_pair_minimal.json"));

  IREE_ASSERT_OK_AND_ASSIGN(
      auto result,
      EncodePair(tokenizer.get(), IREE_SV("hello"), IREE_SV("world")));

  ASSERT_EQ(result.token_ids.size(), 5u);
  EXPECT_EQ(result.token_ids[0], 101);   // [CLS]
  EXPECT_EQ(result.token_ids[1], 7592);  // hello
  EXPECT_EQ(result.token_ids[2], 102);   // [SEP] (infix)
  EXPECT_EQ(result.token_ids[3], 2088);  // world
  EXPECT_EQ(result.token_ids[4], 102);   // [SEP] (suffix)

  ASSERT_EQ(result.type_ids.size(), 5u);
  EXPECT_EQ(result.type_ids[0], 0);  // [CLS] prefix type_id
  EXPECT_EQ(result.type_ids[1], 0);  // sequence A type_id
  EXPECT_EQ(result.type_ids[2], 0);  // [SEP] infix type_id
  EXPECT_EQ(result.type_ids[3], 1);  // sequence B type_id
  EXPECT_EQ(result.type_ids[4], 1);  // [SEP] suffix type_id
}

TEST(PairEncode, BertPairAllowsEmptySecondSequence) {
  // Empty sequence B is still a pair: [CLS] hello [SEP] [SEP].
  IREE_ASSERT_OK_AND_ASSIGN(auto tokenizer,
                            LoadTokenizer("wordpiece_bert_pair_minimal.json"));

  IREE_ASSERT_OK_AND_ASSIGN(
      auto result,
      EncodePair(tokenizer.get(), IREE_SV("hello"), iree_string_view_empty()));

  ASSERT_EQ(result.token_ids.size(), 4u);
  EXPECT_EQ(result.token_ids[0], 101);   // [CLS]
  EXPECT_EQ(result.token_ids[1], 7592);  // hello
  EXPECT_EQ(result.token_ids[2], 102);   // [SEP] (infix)
  EXPECT_EQ(result.token_ids[3], 102);   // [SEP] (suffix)

  ASSERT_EQ(result.type_ids.size(), 4u);
  EXPECT_EQ(result.type_ids[0], 0);
  EXPECT_EQ(result.type_ids[1], 0);
  EXPECT_EQ(result.type_ids[2], 0);
  EXPECT_EQ(result.type_ids[3], 1);
}

TEST(PairEncode, BertPairMultiWord) {
  // Multi-word pair encoding: "deep learning" / "the weather today"
  // Expected token_ids: [101, 2784, 4083, 102, 1996, 4633, 2651, 102]
  // Expected type_ids:  [  0,    0,    0,   0,    1,    1,    1,   1]
  IREE_ASSERT_OK_AND_ASSIGN(auto tokenizer,
                            LoadTokenizer("wordpiece_bert_pair_minimal.json"));

  IREE_ASSERT_OK_AND_ASSIGN(
      auto result, EncodePair(tokenizer.get(), IREE_SV("deep learning"),
                              IREE_SV("the weather today")));

  ASSERT_EQ(result.token_ids.size(), 8u);
  EXPECT_EQ(result.token_ids[0], 101);   // [CLS]
  EXPECT_EQ(result.token_ids[1], 2784);  // deep
  EXPECT_EQ(result.token_ids[2], 4083);  // learning
  EXPECT_EQ(result.token_ids[3], 102);   // [SEP]
  EXPECT_EQ(result.token_ids[4], 1996);  // the
  EXPECT_EQ(result.token_ids[5], 4633);  // weather
  EXPECT_EQ(result.token_ids[6], 2651);  // today
  EXPECT_EQ(result.token_ids[7], 102);   // [SEP]

  // Type IDs: 0 for CLS+query+infix-SEP, 1 for document+suffix-SEP.
  for (size_t i = 0; i <= 3; ++i) {
    EXPECT_EQ(result.type_ids[i], 0) << "type_ids[" << i << "]";
  }
  for (size_t i = 4; i <= 7; ++i) {
    EXPECT_EQ(result.type_ids[i], 1) << "type_ids[" << i << "]";
  }
}

TEST(PairEncode, BatchMixedSingleAndPair) {
  // Batch with one single-sequence item and one pair item.
  IREE_ASSERT_OK_AND_ASSIGN(auto tokenizer,
                            LoadTokenizer("wordpiece_bert_pair_minimal.json"));

  iree_host_size_t state_size = 0;
  IREE_ASSERT_OK(
      iree_tokenizer_encode_state_calculate_size(tokenizer.get(), &state_size));
  std::vector<uint8_t> state_storage(state_size);
  std::vector<uint8_t> transform_buffer(4096);

  // Item 0: single-sequence "hello"
  std::vector<iree_tokenizer_token_id_t> ids0(256);
  std::vector<uint8_t> type_ids0(256, 0xFF);
  iree_tokenizer_encode_batch_item_t items[2] = {};
  items[0].text = IREE_SV("hello");
  items[0].output = iree_tokenizer_make_token_output(ids0.data(), nullptr,
                                                     type_ids0.data(), 256);

  // Item 1: pair "hello" / "world"
  std::vector<iree_tokenizer_token_id_t> ids1(256);
  std::vector<uint8_t> type_ids1(256, 0xFF);
  items[1].text = IREE_SV("hello");
  items[1].text_pair = IREE_SV("world");
  items[1].flags = IREE_TOKENIZER_ENCODE_BATCH_ITEM_FLAG_HAS_TEXT_PAIR;
  items[1].output = iree_tokenizer_make_token_output(ids1.data(), nullptr,
                                                     type_ids1.data(), 256);

  IREE_ASSERT_OK(iree_tokenizer_encode_batch(
      tokenizer.get(), items, 2, IREE_TOKENIZER_ENCODE_FLAG_ADD_SPECIAL_TOKENS,
      iree_make_byte_span(state_storage.data(), state_storage.size()),
      iree_make_byte_span(transform_buffer.data(), transform_buffer.size()),
      iree_tokenizer_offset_run_list_empty()));

  // Item 0: [CLS] hello [SEP]
  ASSERT_EQ(items[0].out_token_count, 3u);
  EXPECT_EQ(ids0[0], 101);   // [CLS]
  EXPECT_EQ(ids0[1], 7592);  // hello
  EXPECT_EQ(ids0[2], 102);   // [SEP]
  EXPECT_EQ(type_ids0[0], 0);
  EXPECT_EQ(type_ids0[1], 0);
  EXPECT_EQ(type_ids0[2], 0);

  // Item 1: [CLS] hello [SEP] world [SEP]
  ASSERT_EQ(items[1].out_token_count, 5u);
  EXPECT_EQ(ids1[0], 101);   // [CLS]
  EXPECT_EQ(ids1[1], 7592);  // hello
  EXPECT_EQ(ids1[2], 102);   // [SEP]
  EXPECT_EQ(ids1[3], 2088);  // world
  EXPECT_EQ(ids1[4], 102);   // [SEP]
  EXPECT_EQ(type_ids1[0], 0);
  EXPECT_EQ(type_ids1[1], 0);
  EXPECT_EQ(type_ids1[2], 0);
  EXPECT_EQ(type_ids1[3], 1);
  EXPECT_EQ(type_ids1[4], 1);
}

TEST(PairEncode, BatchPairThenSingleResetsPostprocessor) {
  IREE_ASSERT_OK_AND_ASSIGN(auto tokenizer,
                            LoadTokenizer("wordpiece_bert_pair_minimal.json"));

  iree_host_size_t state_size = 0;
  IREE_ASSERT_OK(
      iree_tokenizer_encode_state_calculate_size(tokenizer.get(), &state_size));
  std::vector<uint8_t> state_storage(state_size);
  std::vector<uint8_t> transform_buffer(4096);

  std::vector<iree_tokenizer_token_id_t> ids0(256);
  std::vector<uint8_t> type_ids0(256, 0xFF);
  std::vector<iree_tokenizer_token_id_t> ids1(256);
  std::vector<uint8_t> type_ids1(256, 0xFF);

  iree_tokenizer_encode_batch_item_t items[2] = {};
  items[0].text = IREE_SV("hello");
  items[0].text_pair = IREE_SV("world");
  items[0].flags = IREE_TOKENIZER_ENCODE_BATCH_ITEM_FLAG_HAS_TEXT_PAIR;
  items[0].output = iree_tokenizer_make_token_output(ids0.data(), nullptr,
                                                     type_ids0.data(), 256);
  items[1].text = IREE_SV("world");
  items[1].output = iree_tokenizer_make_token_output(ids1.data(), nullptr,
                                                     type_ids1.data(), 256);

  IREE_ASSERT_OK(iree_tokenizer_encode_batch(
      tokenizer.get(), items, 2, IREE_TOKENIZER_ENCODE_FLAG_ADD_SPECIAL_TOKENS,
      iree_make_byte_span(state_storage.data(), state_storage.size()),
      iree_make_byte_span(transform_buffer.data(), transform_buffer.size()),
      iree_tokenizer_offset_run_list_empty()));

  ASSERT_EQ(items[0].out_token_count, 5u);
  EXPECT_EQ(ids0[0], 101);
  EXPECT_EQ(ids0[1], 7592);
  EXPECT_EQ(ids0[2], 102);
  EXPECT_EQ(ids0[3], 2088);
  EXPECT_EQ(ids0[4], 102);
  EXPECT_EQ(type_ids0[3], 1);
  EXPECT_EQ(type_ids0[4], 1);

  ASSERT_EQ(items[1].out_token_count, 3u);
  EXPECT_EQ(ids1[0], 101);
  EXPECT_EQ(ids1[1], 2088);
  EXPECT_EQ(ids1[2], 102);
  EXPECT_EQ(type_ids1[0], 0);
  EXPECT_EQ(type_ids1[1], 0);
  EXPECT_EQ(type_ids1[2], 0);
}

TEST(PairEncode, PairWithoutAddSpecialTokens) {
  // Without ADD_SPECIAL_TOKENS, pair encoding should produce just the
  // concatenated model tokens (no [CLS], [SEP], or type_ids assignment).
  IREE_ASSERT_OK_AND_ASSIGN(auto tokenizer,
                            LoadTokenizer("wordpiece_bert_pair_minimal.json"));

  iree_host_size_t state_size = 0;
  IREE_ASSERT_OK(
      iree_tokenizer_encode_state_calculate_size(tokenizer.get(), &state_size));
  std::vector<uint8_t> state_storage(state_size);
  std::vector<uint8_t> transform_buffer(4096);

  std::vector<iree_tokenizer_token_id_t> token_ids(256);
  std::vector<uint8_t> type_ids(256, 0xFF);

  iree_tokenizer_encode_batch_item_t item = {};
  item.text = IREE_SV("hello");
  item.text_pair = IREE_SV("world");
  item.flags = IREE_TOKENIZER_ENCODE_BATCH_ITEM_FLAG_HAS_TEXT_PAIR;
  item.output = iree_tokenizer_make_token_output(token_ids.data(), nullptr,
                                                 type_ids.data(), 256);

  IREE_ASSERT_OK(iree_tokenizer_encode_batch(
      tokenizer.get(), &item, 1, IREE_TOKENIZER_ENCODE_FLAG_NONE,
      iree_make_byte_span(state_storage.data(), state_storage.size()),
      iree_make_byte_span(transform_buffer.data(), transform_buffer.size()),
      iree_tokenizer_offset_run_list_empty()));

  // Without special tokens: just "hello" + "world" (no [CLS]/[SEP]).
  // Type IDs should all be 0xFF (untouched) since the postprocessor is IDLE.
  ASSERT_EQ(item.out_token_count, 2u);
  EXPECT_EQ(token_ids[0], 7592);  // hello
  EXPECT_EQ(token_ids[1], 2088);  // world
  EXPECT_EQ(type_ids[0], 0xFF);   // Untouched (IDLE postprocessor).
  EXPECT_EQ(type_ids[1], 0xFF);
}

TEST(PairEncode, TextPairRequiresFlag) {
  IREE_ASSERT_OK_AND_ASSIGN(auto tokenizer,
                            LoadTokenizer("wordpiece_bert_pair_minimal.json"));

  iree_host_size_t state_size = 0;
  IREE_ASSERT_OK(
      iree_tokenizer_encode_state_calculate_size(tokenizer.get(), &state_size));
  std::vector<uint8_t> state_storage(state_size);
  std::vector<uint8_t> transform_buffer(4096);
  std::vector<iree_tokenizer_token_id_t> token_ids(16);

  iree_tokenizer_encode_batch_item_t item = {};
  item.text = IREE_SV("hello");
  item.text_pair = IREE_SV("world");
  item.output = iree_tokenizer_make_token_output(token_ids.data(), nullptr,
                                                 nullptr, token_ids.size());

  iree_status_t status = iree_tokenizer_encode_batch(
      tokenizer.get(), &item, 1, IREE_TOKENIZER_ENCODE_FLAG_ADD_SPECIAL_TOKENS,
      iree_make_byte_span(state_storage.data(), state_storage.size()),
      iree_make_byte_span(transform_buffer.data(), transform_buffer.size()),
      iree_tokenizer_offset_run_list_empty());
  IREE_EXPECT_STATUS_IS(IREE_STATUS_INVALID_ARGUMENT, status);
}

TEST(PairEncode, UnknownItemFlagsRejected) {
  IREE_ASSERT_OK_AND_ASSIGN(auto tokenizer,
                            LoadTokenizer("wordpiece_bert_pair_minimal.json"));

  iree_host_size_t state_size = 0;
  IREE_ASSERT_OK(
      iree_tokenizer_encode_state_calculate_size(tokenizer.get(), &state_size));
  std::vector<uint8_t> state_storage(state_size);
  std::vector<uint8_t> transform_buffer(4096);
  std::vector<iree_tokenizer_token_id_t> token_ids(16);

  iree_tokenizer_encode_batch_item_t item = {};
  item.text = IREE_SV("hello");
  item.flags = 0x80000000u;
  item.output = iree_tokenizer_make_token_output(token_ids.data(), nullptr,
                                                 nullptr, token_ids.size());

  iree_status_t status = iree_tokenizer_encode_batch(
      tokenizer.get(), &item, 1, IREE_TOKENIZER_ENCODE_FLAG_ADD_SPECIAL_TOKENS,
      iree_make_byte_span(state_storage.data(), state_storage.size()),
      iree_make_byte_span(transform_buffer.data(), transform_buffer.size()),
      iree_tokenizer_offset_run_list_empty());
  IREE_EXPECT_STATUS_IS(IREE_STATUS_INVALID_ARGUMENT, status);
}

TEST(PairEncode, InvalidStateStorageReturnsStatus) {
  IREE_ASSERT_OK_AND_ASSIGN(auto tokenizer,
                            LoadTokenizer("wordpiece_bert_pair_minimal.json"));

  uint8_t state_storage = 0;
  std::vector<uint8_t> transform_buffer(4096);
  std::vector<iree_tokenizer_token_id_t> token_ids(16);

  iree_tokenizer_encode_batch_item_t item = {};
  item.text = IREE_SV("hello");
  item.text_pair = IREE_SV("world");
  item.flags = IREE_TOKENIZER_ENCODE_BATCH_ITEM_FLAG_HAS_TEXT_PAIR;
  item.output = iree_tokenizer_make_token_output(token_ids.data(), nullptr,
                                                 nullptr, token_ids.size());

  iree_status_t status = iree_tokenizer_encode_batch(
      tokenizer.get(), &item, 1, IREE_TOKENIZER_ENCODE_FLAG_ADD_SPECIAL_TOKENS,
      iree_make_byte_span(&state_storage, 1),
      iree_make_byte_span(transform_buffer.data(), transform_buffer.size()),
      iree_tokenizer_offset_run_list_empty());
  IREE_EXPECT_STATUS_IS(IREE_STATUS_INVALID_ARGUMENT, status);
}

TEST(PairEncode, OutputBufferTooSmallAtPairBoundaries) {
  IREE_ASSERT_OK_AND_ASSIGN(auto tokenizer,
                            LoadTokenizer("wordpiece_bert_pair_minimal.json"));

  for (iree_host_size_t capacity = 0; capacity < 5; ++capacity) {
    auto result = EncodePairDetailed(tokenizer.get(), IREE_SV("hello"),
                                     IREE_SV("world"), capacity,
                                     /*track_offsets=*/false);
    IREE_EXPECT_STATUS_IS(IREE_STATUS_RESOURCE_EXHAUSTED, result.status())
        << "capacity=" << capacity;
  }

  IREE_ASSERT_OK_AND_ASSIGN(
      auto result, EncodePairDetailed(tokenizer.get(), IREE_SV("hello"),
                                      IREE_SV("world"), /*output_capacity=*/5,
                                      /*track_offsets=*/false));
  EXPECT_EQ(result.token_ids, (std::vector<iree_tokenizer_token_id_t>{
                                  101, 7592, 102, 2088, 102}));
}

TEST(PairEncode, OffsetsAreRelativeToEachSequence) {
  IREE_ASSERT_OK_AND_ASSIGN(auto tokenizer,
                            LoadTokenizer("wordpiece_bert_pair_minimal.json"));

  IREE_ASSERT_OK_AND_ASSIGN(
      auto result, EncodePairDetailed(tokenizer.get(), IREE_SV("hello"),
                                      IREE_SV("world"), /*output_capacity=*/8,
                                      /*track_offsets=*/true));

  ASSERT_EQ(result.token_offsets.size(), 5u);
  EXPECT_EQ(result.token_offsets[0].start, 0u);
  EXPECT_EQ(result.token_offsets[0].end, 0u);
  EXPECT_EQ(result.token_offsets[1].start, 0u);
  EXPECT_EQ(result.token_offsets[1].end, 5u);
  EXPECT_EQ(result.token_offsets[2].start, 0u);
  EXPECT_EQ(result.token_offsets[2].end, 0u);
  EXPECT_EQ(result.token_offsets[3].start, 0u);
  EXPECT_EQ(result.token_offsets[3].end, 5u);
  EXPECT_EQ(result.token_offsets[4].start, 0u);
  EXPECT_EQ(result.token_offsets[4].end, 0u);
}

}  // namespace
}  // namespace iree::tokenizer
