// 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

// Tokenizes text using HuggingFace tokenizer.json files.
//
// Example encoding text to token IDs (default: comma-separated):
//   iree-tokenize --tokenizer=tokenizer.json "Hello, world!"
//   # Output: 101,7592,1010,2088,999,102
//
// Example JSON output:
//   iree-tokenize --tokenizer=tokenizer.json --json "Hello, world!"
//   # Output: {"ids":[101,7592,1010,2088,999,102]}
//
// Example encoding without special tokens:
//   iree-tokenize --tokenizer=tokenizer.json --special=false "Hello, world!"
//   # Output: 7592,1010,2088,999
//
// Example with offset tracking:
//   iree-tokenize --tokenizer=tokenizer.json --offsets "Hello, world!"
//   # Output: 7592[0:5],1010[5:6],2088[7:12],999[12:13]
//
// Example decoding token IDs to text:
//   iree-tokenize --tokenizer=tokenizer.json --decode
//   "101,7592,1010,2088,999,102"
//   # Output: Hello, world!
//
// Example batch mode (one line per input):
//   echo -e "Hello\nWorld" | iree-tokenize --tokenizer=tokenizer.json --batch
//   # Output: 101,7592,...
//   #         101,2088,...
//
// Example showing tokenizer info (always JSON):
//   iree-tokenize --tokenizer=tokenizer.json --info
//   # Output: {"vocab_size":30522,"model_type":"BPE",...}
//
// Example benchmarking:
//   iree-tokenize --tokenizer=tokenizer.json --benchmark=oneshot "Hello,
//   world!" # Output: timing stats to stderr, token IDs to stdout

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "iree/base/api.h"
#include "iree/base/internal/json.h"
#include "iree/base/tooling/flags.h"
#include "iree/io/file_contents.h"
#include "iree/tokenizer/format/huggingface/tokenizer_json.h"
#include "iree/tokenizer/format/tiktoken/tiktoken.h"
#include "iree/tokenizer/tokenizer.h"
#include "iree/tokenizer/vocab/vocab.h"

//===----------------------------------------------------------------------===//
// Flags
//===----------------------------------------------------------------------===//

IREE_FLAG(bool, decode, false, "Decode mode: input is comma-separated IDs.");
IREE_FLAG(bool, decode_special, false,
          "Include special tokens (BOS/EOS) in decode output.");
IREE_FLAG(bool, special, true, "Add special tokens (BOS/EOS, CLS/SEP).");
IREE_FLAG(bool, batch, false, "Batch mode: read lines from stdin.");
IREE_FLAG(bool, stream, false, "Stream stdin continuously (not line-by-line).");
IREE_FLAG(int32_t, max_length, 0, "Max output length (0 = unlimited).");
IREE_FLAG(bool, info, false, "Show tokenizer info instead of encoding.");
IREE_FLAG(bool, json, false,
          "Output JSON format (default: comma-separated IDs).");
IREE_FLAG(bool, json_string, false,
          "Input is a JSON-encoded string (handles \\uXXXX escapes).");
IREE_FLAG(string, tokenizer, "",
          "Path to tokenizer file (.json for HuggingFace, .tiktoken for "
          "OpenAI tiktoken format).");
IREE_FLAG(string, encoding, "",
          "Tiktoken encoding name (cl100k_base, o200k_base, r50k_base, "
          "p50k_base). Required for .tiktoken files; ignored for .json.");
IREE_FLAG(bool, offsets, false, "Show token-to-byte offset mappings.");
IREE_FLAG(string, benchmark, "",
          "Benchmark mode: oneshot, batch, stream, or decode.");
IREE_FLAG(int32_t, benchmark_iterations, 100,
          "Number of timed iterations for benchmarking.");
IREE_FLAG(int32_t, benchmark_warmup, 5,
          "Number of warmup iterations before timing.");
IREE_FLAG(int32_t, benchmark_chunk_size, 4096,
          "Chunk size in bytes for stream benchmark.");
//===----------------------------------------------------------------------===//
// Output Helpers
//===----------------------------------------------------------------------===//

// Prints a token ID with optional offset annotation.
static void iree_tooling_print_token(iree_tokenizer_token_id_t token_id,
                                     const iree_tokenizer_offset_t* offset,
                                     bool first) {
  if (!first) fputc(',', stdout);
  fprintf(stdout, "%" PRId32, token_id);
  if (offset) {
    fprintf(stdout, "[%zu:%zu]", (size_t)offset->start, (size_t)offset->end);
  }
}

// Prints a token sequence with optional offsets.
static void iree_tooling_print_tokens(
    const iree_tokenizer_token_id_t* token_ids,
    const iree_tokenizer_offset_t* offsets, iree_host_size_t count,
    bool* first) {
  for (iree_host_size_t i = 0; i < count; ++i) {
    iree_tooling_print_token(token_ids[i], offsets ? &offsets[i] : NULL,
                             *first);
    *first = false;
  }
}

// Prints a JSON array of token IDs with optional offsets.
static void iree_tooling_print_json_tokens(
    const iree_tokenizer_token_id_t* token_ids,
    const iree_tokenizer_offset_t* offsets, iree_host_size_t count) {
  fputs("{\"ids\":[", stdout);
  for (iree_host_size_t i = 0; i < count; ++i) {
    if (i > 0) fputc(',', stdout);
    fprintf(stdout, "%" PRId32, token_ids[i]);
  }
  fputc(']', stdout);
  if (offsets) {
    fputs(",\"offsets\":[", stdout);
    for (iree_host_size_t i = 0; i < count; ++i) {
      if (i > 0) fputc(',', stdout);
      fprintf(stdout, "[%zu,%zu]", (size_t)offsets[i].start,
              (size_t)offsets[i].end);
    }
    fputc(']', stdout);
  }
  fputs("}\n", stdout);
}

// Writes decoded text to stdout with JSON escaping if needed.
static void iree_tooling_print_text(const char* data, iree_host_size_t length,
                                    bool json_escape) {
  if (!json_escape) {
    fwrite(data, 1, length, stdout);
    return;
  }
  for (iree_host_size_t i = 0; i < length; ++i) {
    char c = data[i];
    switch (c) {
      case '"':
        fputs("\\\"", stdout);
        break;
      case '\\':
        fputs("\\\\", stdout);
        break;
      case '\b':
        fputs("\\b", stdout);
        break;
      case '\f':
        fputs("\\f", stdout);
        break;
      case '\n':
        fputs("\\n", stdout);
        break;
      case '\r':
        fputs("\\r", stdout);
        break;
      case '\t':
        fputs("\\t", stdout);
        break;
      default:
        if ((unsigned char)c < 0x20) {
          fprintf(stdout, "\\u%04x", (unsigned char)c);
        } else {
          fputc(c, stdout);
        }
        break;
    }
  }
}

//===----------------------------------------------------------------------===//
// Encode (one-shot with retry)
//===----------------------------------------------------------------------===//

// Builds encode flags from CLI flags.
static iree_tokenizer_encode_flags_t iree_tooling_encode_flags(void) {
  iree_tokenizer_encode_flags_t flags =
      IREE_TOKENIZER_ENCODE_FLAG_AT_INPUT_START;
  if (FLAG_special) flags |= IREE_TOKENIZER_ENCODE_FLAG_ADD_SPECIAL_TOKENS;
  if (FLAG_offsets) flags |= IREE_TOKENIZER_ENCODE_FLAG_TRACK_OFFSETS;
  return flags;
}

static iree_status_t iree_tooling_tokenize_encode(
    const iree_tokenizer_t* tokenizer, iree_string_view_t text,
    iree_allocator_t allocator) {
  IREE_TRACE_ZONE_BEGIN(z0);

  iree_tokenizer_encode_flags_t flags = iree_tooling_encode_flags();

  // Allocate combined output buffer. Use fixed capacity - the streaming encode
  // API operates in bounded memory and should never need retries.
  iree_host_size_t capacity = 8192;
  iree_host_size_t total_size = 0;
  iree_host_size_t token_ids_offset = 0;
  iree_host_size_t offsets_offset = 0;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, IREE_STRUCT_LAYOUT(
              0, &total_size,
              IREE_STRUCT_FIELD(capacity, iree_tokenizer_token_id_t,
                                &token_ids_offset),
              IREE_STRUCT_FIELD(FLAG_offsets ? capacity : 0,
                                iree_tokenizer_offset_t, &offsets_offset)));
  uint8_t* storage = NULL;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_allocator_malloc(allocator, total_size, (void**)&storage));

  iree_tokenizer_token_id_t* token_ids =
      (iree_tokenizer_token_id_t*)(storage + token_ids_offset);
  iree_tokenizer_offset_t* offsets =
      FLAG_offsets ? (iree_tokenizer_offset_t*)(storage + offsets_offset)
                   : NULL;
  iree_host_size_t token_count = 0;

  iree_tokenizer_token_output_t output =
      iree_tokenizer_make_token_output(token_ids, offsets, NULL, capacity);
  iree_status_t status = iree_tokenizer_encode(tokenizer, text, flags, output,
                                               allocator, &token_count);

  if (iree_status_is_ok(status)) {
    // Apply max_length truncation.
    if (FLAG_max_length > 0 &&
        token_count > (iree_host_size_t)FLAG_max_length) {
      token_count = (iree_host_size_t)FLAG_max_length;
    }

    if (FLAG_json) {
      iree_tooling_print_json_tokens(token_ids, offsets, token_count);
    } else {
      bool first = true;
      iree_tooling_print_tokens(token_ids, offsets, token_count, &first);
      fputc('\n', stdout);
    }
  }

  iree_allocator_free(allocator, storage);

  IREE_TRACE_ZONE_END(z0);
  return status;
}

//===----------------------------------------------------------------------===//
// Decode (one-shot with retry)
//===----------------------------------------------------------------------===//

// Parses comma-separated IDs from a string.
static iree_status_t iree_tooling_parse_ids(iree_string_view_t text,
                                            iree_tokenizer_token_id_t* out_ids,
                                            iree_host_size_t max_ids,
                                            iree_host_size_t* out_count) {
  *out_count = 0;
  if (text.size == 0) return iree_ok_status();

  iree_host_size_t position = 0;
  while (position < text.size) {
    // Skip whitespace.
    while (position < text.size &&
           (text.data[position] == ' ' || text.data[position] == '\t')) {
      ++position;
    }
    if (position >= text.size) break;

    // Parse number.
    bool negative = false;
    if (text.data[position] == '-') {
      negative = true;
      ++position;
    }
    int32_t value = 0;
    bool found_digit = false;
    while (position < text.size && text.data[position] >= '0' &&
           text.data[position] <= '9') {
      value = value * 10 + (text.data[position] - '0');
      found_digit = true;
      ++position;
    }
    if (!found_digit) {
      return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                              "expected number at position %zu", position);
    }
    if (negative) value = -value;

    if (*out_count >= max_ids) {
      return iree_make_status(IREE_STATUS_RESOURCE_EXHAUSTED,
                              "too many IDs (max %zu)", max_ids);
    }
    out_ids[(*out_count)++] = value;

    // Skip comma.
    while (position < text.size &&
           (text.data[position] == ' ' || text.data[position] == '\t')) {
      ++position;
    }
    if (position < text.size && text.data[position] == ',') {
      ++position;
    }
  }
  return iree_ok_status();
}

static iree_status_t iree_tooling_tokenize_decode(
    const iree_tokenizer_t* tokenizer, iree_string_view_t input,
    iree_allocator_t allocator) {
  IREE_TRACE_ZONE_BEGIN(z0);

  // Parse IDs into stack buffer.
  iree_tokenizer_token_id_t ids[8192];
  iree_host_size_t id_count = 0;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_tooling_parse_ids(input, ids, IREE_ARRAYSIZE(ids), &id_count));

  iree_tokenizer_token_id_list_t tokens =
      iree_tokenizer_make_token_id_list(ids, id_count);

  // Decode with retry on RESOURCE_EXHAUSTED.
  iree_host_size_t text_capacity = 65536;
  char* text_buffer = NULL;
  iree_host_size_t text_length = 0;
  iree_status_t status = iree_ok_status();

  for (;;) {
    status =
        iree_allocator_malloc(allocator, text_capacity, (void**)&text_buffer);
    if (!iree_status_is_ok(status)) break;

    iree_mutable_string_view_t text_output = {text_buffer, text_capacity};
    iree_tokenizer_decode_flags_t decode_flags =
        FLAG_decode_special ? IREE_TOKENIZER_DECODE_FLAG_NONE
                            : IREE_TOKENIZER_DECODE_FLAG_SKIP_SPECIAL_TOKENS;
    status = iree_tokenizer_decode(tokenizer, tokens, decode_flags, text_output,
                                   allocator, &text_length);
    if (iree_status_is_resource_exhausted(status)) {
      iree_status_ignore(status);
      iree_allocator_free(allocator, text_buffer);
      text_buffer = NULL;
      text_capacity *= 2;
      continue;
    }
    break;
  }

  if (iree_status_is_ok(status)) {
    if (FLAG_json) {
      fputs("{\"text\":\"", stdout);
      iree_tooling_print_text(text_buffer, text_length, /*json_escape=*/true);
      fputs("\"}\n", stdout);
    } else {
      iree_tooling_print_text(text_buffer, text_length, /*json_escape=*/false);
      fputc('\n', stdout);
    }
  }

  iree_allocator_free(allocator, text_buffer);
  IREE_TRACE_ZONE_END(z0);
  return status;
}

//===----------------------------------------------------------------------===//
// Info
//===----------------------------------------------------------------------===//

static iree_status_t iree_tooling_tokenize_info(
    const iree_tokenizer_t* tokenizer) {
  IREE_TRACE_ZONE_BEGIN(z0);

  const iree_tokenizer_vocab_t* vocab = iree_tokenizer_vocab(tokenizer);
  iree_host_size_t vocab_size = iree_tokenizer_vocab_capacity(vocab);
  iree_host_size_t merge_count = iree_tokenizer_vocab_merge_count(vocab);
  iree_tokenizer_special_ids_t special =
      iree_tokenizer_vocab_special_ids(vocab);
  iree_string_view_t model_type = iree_tokenizer_model_type_name(tokenizer);

  fprintf(stdout, "{\"vocab_size\":%zu,\"model_type\":\"%.*s\"",
          (size_t)vocab_size, (int)model_type.size, model_type.data);

  if (merge_count > 0) {
    fprintf(stdout, ",\"merge_count\":%zu", (size_t)merge_count);
  }

  // Special tokens.
  if (special.bos >= 0) fprintf(stdout, ",\"bos_id\":%" PRId32, special.bos);
  if (special.eos >= 0) fprintf(stdout, ",\"eos_id\":%" PRId32, special.eos);
  if (special.unk >= 0) fprintf(stdout, ",\"unk_id\":%" PRId32, special.unk);
  if (special.pad >= 0) fprintf(stdout, ",\"pad_id\":%" PRId32, special.pad);
  if (special.cls >= 0) fprintf(stdout, ",\"cls_id\":%" PRId32, special.cls);
  if (special.sep >= 0) fprintf(stdout, ",\"sep_id\":%" PRId32, special.sep);
  if (special.mask >= 0) fprintf(stdout, ",\"mask_id\":%" PRId32, special.mask);

  fputs("}\n", stdout);

  IREE_TRACE_ZONE_END(z0);
  return iree_ok_status();
}

//===----------------------------------------------------------------------===//
// Streaming Stdin Mode
//===----------------------------------------------------------------------===//

// Streams stdin, reading chunks and emitting tokens incrementally.
// Uses the pull-based streaming encode API which handles all boundary
// conditions:
// - Incomplete UTF-8 sequences at chunk boundaries
// - Literals (added_tokens) that span chunks
// - Transform segments that span chunks
// - BOS/EOS token emission (via postprocessor)
static iree_status_t iree_tooling_tokenize_stdin_streaming(
    const iree_tokenizer_t* tokenizer, iree_allocator_t allocator) {
  IREE_TRACE_ZONE_BEGIN(z0);

  iree_tokenizer_encode_flags_t flags = iree_tooling_encode_flags();

  // Calculate state storage requirements.
  iree_host_size_t state_size = 0;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_tokenizer_encode_state_calculate_size(tokenizer, &state_size));

  // Allocate combined state and transform buffer.
  iree_host_size_t transform_size =
      iree_tokenizer_transform_buffer_recommended_size(8192);
  iree_host_size_t total_size = 0;
  iree_host_size_t state_offset = 0;
  iree_host_size_t transform_offset = 0;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0,
      IREE_STRUCT_LAYOUT(
          0, &total_size, IREE_STRUCT_FIELD(state_size, uint8_t, &state_offset),
          IREE_STRUCT_FIELD(transform_size, uint8_t, &transform_offset)));
  uint8_t* storage = NULL;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_allocator_malloc(allocator, total_size, (void**)&storage));

  iree_byte_span_t state_span = {storage + state_offset, state_size};
  iree_byte_span_t transform_span = {storage + transform_offset,
                                     transform_size};

  // Initialize streaming state.
  iree_tokenizer_encode_state_t* state = NULL;
  iree_status_t status = iree_tokenizer_encode_state_initialize(
      tokenizer, state_span, transform_span,
      iree_tokenizer_offset_run_list_empty(), flags, &state);

  if (!iree_status_is_ok(status)) {
    iree_allocator_free(allocator, storage);
    IREE_TRACE_ZONE_END(z0);
    return status;
  }

  // Token output buffer (reused each feed call).
  iree_tokenizer_token_id_t token_buffer[1024];
  iree_tokenizer_token_output_t output = iree_tokenizer_make_token_output(
      token_buffer, NULL, NULL, IREE_ARRAYSIZE(token_buffer));

  // Start output.
  if (FLAG_json) fputs("{\"ids\":[", stdout);
  bool first_token = true;

  // Read and feed chunks until EOF.
  char read_buffer[8192];
  size_t bytes_read;
  while (iree_status_is_ok(status) &&
         (bytes_read = fread(read_buffer, 1, sizeof(read_buffer), stdin)) > 0) {
    iree_string_view_t chunk = iree_make_string_view(read_buffer, bytes_read);
    while (chunk.size > 0 && iree_status_is_ok(status)) {
      iree_host_size_t bytes_consumed = 0;
      iree_host_size_t token_count = 0;
      status = iree_tokenizer_encode_state_feed(state, chunk, output,
                                                &bytes_consumed, &token_count);
      if (iree_status_is_ok(status)) {
        iree_tooling_print_tokens(token_buffer, NULL, token_count,
                                  &first_token);
        chunk.data += bytes_consumed;
        chunk.size -= bytes_consumed;
      }
    }
  }

  // Finalize: flush any pending state.
  if (iree_status_is_ok(status)) {
    iree_host_size_t token_count = 0;
    status = iree_tokenizer_encode_state_finalize(state, output, &token_count);
    if (iree_status_is_ok(status)) {
      iree_tooling_print_tokens(token_buffer, NULL, token_count, &first_token);
    }
  }

  // Close output.
  if (FLAG_json) {
    fputs("]}\n", stdout);
  } else {
    fputc('\n', stdout);
  }

  iree_tokenizer_encode_state_deinitialize(state);
  iree_allocator_free(allocator, storage);

  IREE_TRACE_ZONE_END(z0);
  return status;
}

//===----------------------------------------------------------------------===//
// Batch Mode (line-by-line)
//===----------------------------------------------------------------------===//

// Strips the trailing LF line terminator from a string view.
// Only strips \n - does NOT strip \r, which could be content.
// The batch protocol uses \n as delimiter (via Python's "\n".join()), so any
// \r before the \n is content that must be preserved.
static iree_string_view_t iree_tooling_string_view_strip_trailing_newline(
    iree_string_view_t text) {
  if (text.size > 0 && text.data[text.size - 1] == '\n') {
    --text.size;
  }
  return text;
}

// Portable getline implementation that dynamically grows the buffer.
// Returns the line length (excluding null terminator), or -1 on EOF/error.
// The caller must free *line_ptr using the same allocator when done.
static intptr_t iree_tooling_getline(char** line_ptr,
                                     iree_host_size_t* capacity_ptr,
                                     FILE* stream, iree_allocator_t allocator) {
  if (*line_ptr == NULL || *capacity_ptr == 0) {
    *capacity_ptr = 256;
    iree_status_t status =
        iree_allocator_malloc(allocator, *capacity_ptr, (void**)line_ptr);
    if (!iree_status_is_ok(status)) {
      iree_status_ignore(status);
      return -1;
    }
  }

  iree_host_size_t position = 0;
  int character;
  while ((character = fgetc(stream)) != EOF) {
    // Grow buffer if needed (leaving room for null terminator).
    if (position + 1 >= *capacity_ptr) {
      iree_host_size_t new_capacity = *capacity_ptr * 2;
      iree_status_t status =
          iree_allocator_realloc(allocator, new_capacity, (void**)line_ptr);
      if (!iree_status_is_ok(status)) {
        iree_status_ignore(status);
        return -1;
      }
      *capacity_ptr = new_capacity;
    }

    (*line_ptr)[position++] = (char)character;
    if (character == '\n') break;
  }

  if (position == 0 && character == EOF) return -1;

  (*line_ptr)[position] = '\0';
  return (intptr_t)position;
}

static iree_status_t iree_tooling_tokenize_batch(
    const iree_tokenizer_t* tokenizer, iree_allocator_t allocator) {
  IREE_TRACE_ZONE_BEGIN(z0);

  char* line = NULL;
  iree_host_size_t line_capacity = 0;
  intptr_t line_length;

  while ((line_length = iree_tooling_getline(&line, &line_capacity, stdin,
                                             allocator)) != -1) {
    iree_string_view_t text = iree_tooling_string_view_strip_trailing_newline(
        iree_make_string_view(line, (iree_host_size_t)line_length));
    iree_status_t status;
    if (FLAG_decode) {
      status = iree_tooling_tokenize_decode(tokenizer, text, allocator);
    } else {
      status = iree_tooling_tokenize_encode(tokenizer, text, allocator);
    }
    if (!iree_status_is_ok(status)) {
      iree_allocator_free(allocator, line);
      IREE_TRACE_ZONE_END(z0);
      return status;
    }
  }

  iree_allocator_free(allocator, line);
  IREE_TRACE_ZONE_END(z0);
  return iree_ok_status();
}

//===----------------------------------------------------------------------===//
// Benchmark Mode
//===----------------------------------------------------------------------===//

typedef struct {
  iree_time_t min_ns;
  iree_time_t max_ns;
  iree_time_t total_ns;
  int32_t iterations;
  iree_host_size_t total_input_bytes;
  iree_host_size_t total_tokens;
  iree_host_size_t peak_memory;
} iree_tooling_benchmark_stats_t;

static void iree_tooling_benchmark_stats_initialize(
    iree_tooling_benchmark_stats_t* stats) {
  memset(stats, 0, sizeof(*stats));
  stats->min_ns = INT64_MAX;
}

static void iree_tooling_benchmark_stats_record(
    iree_tooling_benchmark_stats_t* stats, iree_time_t elapsed_ns,
    iree_host_size_t input_bytes, iree_host_size_t tokens) {
  if (elapsed_ns < stats->min_ns) stats->min_ns = elapsed_ns;
  if (elapsed_ns > stats->max_ns) stats->max_ns = elapsed_ns;
  stats->total_ns += elapsed_ns;
  stats->iterations++;
  stats->total_input_bytes += input_bytes;
  stats->total_tokens += tokens;
}

static void iree_tooling_benchmark_stats_print(
    const iree_tooling_benchmark_stats_t* stats, const char* mode) {
  iree_time_t average_ns = stats->total_ns / stats->iterations;
  double tokens_per_sec =
      (double)stats->total_tokens / ((double)stats->total_ns / 1e9);
  double mb_per_sec =
      (double)stats->total_input_bytes / ((double)stats->total_ns / 1e9) / 1e6;

  if (FLAG_json) {
    fprintf(stdout,
            "{\"mode\":\"%s\",\"iterations\":%d,"
            "\"total_input_bytes\":%zu,\"total_tokens\":%zu,"
            "\"min_ns\":%" PRId64 ",\"avg_ns\":%" PRId64 ",\"max_ns\":%" PRId64
            ","
            "\"tokens_per_sec\":%.1f,\"mb_per_sec\":%.3f,"
            "\"peak_memory_bytes\":%zu}\n",
            mode, stats->iterations, (size_t)stats->total_input_bytes,
            (size_t)stats->total_tokens, stats->min_ns, average_ns,
            stats->max_ns, tokens_per_sec, mb_per_sec,
            (size_t)stats->peak_memory);
  } else {
    fprintf(stderr,
            "Benchmark: %s\n"
            "  Iterations:    %d\n"
            "  Input bytes:   %zu total\n"
            "  Tokens:        %zu total\n"
            "  Latency (ns):  min=%" PRId64 " avg=%" PRId64 " max=%" PRId64
            "\n"
            "  Throughput:    %.1f tokens/sec, %.3f MB/sec\n"
            "  Peak memory:   %zu bytes\n",
            mode, stats->iterations, (size_t)stats->total_input_bytes,
            (size_t)stats->total_tokens, stats->min_ns, average_ns,
            stats->max_ns, tokens_per_sec, mb_per_sec,
            (size_t)stats->peak_memory);
  }
}

static iree_status_t iree_tooling_benchmark_oneshot(
    const iree_tokenizer_t* tokenizer, iree_string_view_t text,
    iree_allocator_t allocator) {
  IREE_TRACE_ZONE_BEGIN(z0);

  iree_tokenizer_encode_flags_t flags = iree_tooling_encode_flags();

  // Allocate output buffer sized to text length (generous).
  iree_host_size_t capacity = iree_max(text.size, (iree_host_size_t)8192);
  iree_tokenizer_token_id_t* token_ids = NULL;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_allocator_malloc(allocator,
                                capacity * sizeof(iree_tokenizer_token_id_t),
                                (void**)&token_ids));

  iree_tokenizer_token_output_t output =
      iree_tokenizer_make_token_output(token_ids, NULL, NULL, capacity);

  iree_tooling_benchmark_stats_t stats;
  iree_tooling_benchmark_stats_initialize(&stats);
  stats.peak_memory = capacity * sizeof(iree_tokenizer_token_id_t);

  // Warmup.
  for (int32_t i = 0; i < FLAG_benchmark_warmup; ++i) {
    iree_host_size_t token_count = 0;
    iree_status_t status = iree_tokenizer_encode(tokenizer, text, flags, output,
                                                 allocator, &token_count);
    if (!iree_status_is_ok(status)) {
      iree_allocator_free(allocator, token_ids);
      IREE_TRACE_ZONE_END(z0);
      return status;
    }
  }

  // Timed iterations.
  for (int32_t i = 0; i < FLAG_benchmark_iterations; ++i) {
    iree_host_size_t token_count = 0;
    iree_time_t start = iree_time_now();
    iree_status_t status = iree_tokenizer_encode(tokenizer, text, flags, output,
                                                 allocator, &token_count);
    iree_time_t end = iree_time_now();
    if (!iree_status_is_ok(status)) {
      iree_allocator_free(allocator, token_ids);
      IREE_TRACE_ZONE_END(z0);
      return status;
    }
    iree_tooling_benchmark_stats_record(&stats, end - start, text.size,
                                        token_count);
  }

  iree_tooling_benchmark_stats_print(&stats, "oneshot");
  iree_allocator_free(allocator, token_ids);
  IREE_TRACE_ZONE_END(z0);
  return iree_ok_status();
}

static iree_status_t iree_tooling_benchmark_stream(
    const iree_tokenizer_t* tokenizer, iree_string_view_t text,
    iree_allocator_t allocator) {
  IREE_TRACE_ZONE_BEGIN(z0);

  iree_tokenizer_encode_flags_t flags = iree_tooling_encode_flags();
  iree_host_size_t chunk_size = (iree_host_size_t)FLAG_benchmark_chunk_size;

  // Allocate combined state and transform buffer.
  iree_host_size_t state_size = 0;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_tokenizer_encode_state_calculate_size(tokenizer, &state_size));
  iree_host_size_t transform_size =
      iree_tokenizer_transform_buffer_recommended_size(chunk_size);
  iree_host_size_t total_size = 0;
  iree_host_size_t state_offset = 0;
  iree_host_size_t transform_offset = 0;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0,
      IREE_STRUCT_LAYOUT(
          0, &total_size, IREE_STRUCT_FIELD(state_size, uint8_t, &state_offset),
          IREE_STRUCT_FIELD(transform_size, uint8_t, &transform_offset)));
  uint8_t* storage = NULL;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_allocator_malloc(allocator, total_size, (void**)&storage));

  // Token output buffer.
  iree_tokenizer_token_id_t token_buffer[1024];
  iree_tokenizer_token_output_t output = iree_tokenizer_make_token_output(
      token_buffer, NULL, NULL, IREE_ARRAYSIZE(token_buffer));

  iree_byte_span_t state_span = {storage + state_offset, state_size};
  iree_byte_span_t transform_span = {storage + transform_offset,
                                     transform_size};

  iree_tooling_benchmark_stats_t stats;
  iree_tooling_benchmark_stats_initialize(&stats);
  stats.peak_memory = total_size + sizeof(token_buffer);

  iree_status_t status = iree_ok_status();
  int32_t total_iterations = FLAG_benchmark_warmup + FLAG_benchmark_iterations;
  for (int32_t iteration = 0;
       iteration < total_iterations && iree_status_is_ok(status); ++iteration) {
    bool is_warmup = (iteration < FLAG_benchmark_warmup);
    iree_time_t start = iree_time_now();
    iree_host_size_t iteration_tokens = 0;

    // Initialize state for this iteration.
    iree_tokenizer_encode_state_t* state = NULL;
    status = iree_tokenizer_encode_state_initialize(
        tokenizer, state_span, transform_span,
        iree_tokenizer_offset_run_list_empty(), flags, &state);
    if (!iree_status_is_ok(status)) break;

    // Feed text in chunks.
    iree_host_size_t text_position = 0;
    while (text_position < text.size && iree_status_is_ok(status)) {
      iree_host_size_t remaining = text.size - text_position;
      iree_host_size_t this_chunk = iree_min(remaining, chunk_size);
      iree_string_view_t chunk =
          iree_make_string_view(text.data + text_position, this_chunk);
      while (chunk.size > 0 && iree_status_is_ok(status)) {
        iree_host_size_t bytes_consumed = 0;
        iree_host_size_t token_count = 0;
        status = iree_tokenizer_encode_state_feed(
            state, chunk, output, &bytes_consumed, &token_count);
        if (iree_status_is_ok(status)) {
          iteration_tokens += token_count;
          chunk.data += bytes_consumed;
          chunk.size -= bytes_consumed;
        }
      }
      text_position += this_chunk;
    }

    // Finalize.
    if (iree_status_is_ok(status)) {
      iree_host_size_t token_count = 0;
      status =
          iree_tokenizer_encode_state_finalize(state, output, &token_count);
      if (iree_status_is_ok(status)) {
        iteration_tokens += token_count;
      }
    }

    iree_tokenizer_encode_state_deinitialize(state);

    if (iree_status_is_ok(status) && !is_warmup) {
      iree_time_t end = iree_time_now();
      iree_tooling_benchmark_stats_record(&stats, end - start, text.size,
                                          iteration_tokens);
    }
  }

  if (iree_status_is_ok(status)) {
    iree_tooling_benchmark_stats_print(&stats, "stream");
  }

  iree_allocator_free(allocator, storage);
  IREE_TRACE_ZONE_END(z0);
  return status;
}

static iree_status_t iree_tooling_benchmark_decode(
    const iree_tokenizer_t* tokenizer, iree_string_view_t text,
    iree_allocator_t allocator) {
  IREE_TRACE_ZONE_BEGIN(z0);

  // First, encode the text to get tokens.
  iree_tokenizer_encode_flags_t flags = iree_tooling_encode_flags();
  iree_host_size_t capacity = iree_max(text.size, (iree_host_size_t)8192);
  iree_tokenizer_token_id_t* token_ids = NULL;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_allocator_malloc(allocator,
                                capacity * sizeof(iree_tokenizer_token_id_t),
                                (void**)&token_ids));

  iree_tokenizer_token_output_t output =
      iree_tokenizer_make_token_output(token_ids, NULL, NULL, capacity);
  iree_host_size_t token_count = 0;
  iree_status_t status = iree_tokenizer_encode(tokenizer, text, flags, output,
                                               allocator, &token_count);
  if (!iree_status_is_ok(status)) {
    iree_allocator_free(allocator, token_ids);
    IREE_TRACE_ZONE_END(z0);
    return status;
  }

  iree_tokenizer_token_id_list_t tokens =
      iree_tokenizer_make_token_id_list(token_ids, token_count);

  // Allocate decode output buffer.
  iree_host_size_t text_capacity = 65536;
  char* text_buffer = NULL;
  status =
      iree_allocator_malloc(allocator, text_capacity, (void**)&text_buffer);
  if (!iree_status_is_ok(status)) {
    iree_allocator_free(allocator, token_ids);
    IREE_TRACE_ZONE_END(z0);
    return status;
  }

  iree_mutable_string_view_t text_output = {text_buffer, text_capacity};

  iree_tooling_benchmark_stats_t stats;
  iree_tooling_benchmark_stats_initialize(&stats);
  stats.peak_memory =
      capacity * sizeof(iree_tokenizer_token_id_t) + text_capacity;

  // Decode flags for benchmark iterations.
  iree_tokenizer_decode_flags_t decode_flags =
      FLAG_decode_special ? IREE_TOKENIZER_DECODE_FLAG_NONE
                          : IREE_TOKENIZER_DECODE_FLAG_SKIP_SPECIAL_TOKENS;

  // Warmup.
  for (int32_t i = 0; i < FLAG_benchmark_warmup && iree_status_is_ok(status);
       ++i) {
    iree_host_size_t text_length = 0;
    status = iree_tokenizer_decode(tokenizer, tokens, decode_flags, text_output,
                                   allocator, &text_length);
  }

  // Timed iterations.
  for (int32_t i = 0;
       i < FLAG_benchmark_iterations && iree_status_is_ok(status); ++i) {
    iree_host_size_t text_length = 0;
    iree_time_t start = iree_time_now();
    status = iree_tokenizer_decode(tokenizer, tokens, decode_flags, text_output,
                                   allocator, &text_length);
    iree_time_t end = iree_time_now();
    if (iree_status_is_ok(status)) {
      iree_tooling_benchmark_stats_record(&stats, end - start, text_length,
                                          token_count);
    }
  }

  if (iree_status_is_ok(status)) {
    iree_tooling_benchmark_stats_print(&stats, "decode");
  }

  iree_allocator_free(allocator, text_buffer);
  iree_allocator_free(allocator, token_ids);
  IREE_TRACE_ZONE_END(z0);
  return status;
}

static iree_status_t iree_tooling_tokenize_benchmark(
    const iree_tokenizer_t* tokenizer, iree_string_view_t text,
    iree_allocator_t allocator) {
  iree_string_view_t mode = iree_make_cstring_view(FLAG_benchmark);

  if (iree_string_view_equal(mode, IREE_SV("oneshot"))) {
    return iree_tooling_benchmark_oneshot(tokenizer, text, allocator);
  } else if (iree_string_view_equal(mode, IREE_SV("stream"))) {
    return iree_tooling_benchmark_stream(tokenizer, text, allocator);
  } else if (iree_string_view_equal(mode, IREE_SV("decode"))) {
    return iree_tooling_benchmark_decode(tokenizer, text, allocator);
  }

  return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                          "unknown benchmark mode '%.*s' "
                          "(expected: oneshot, stream, decode)",
                          (int)mode.size, mode.data);
}

//===----------------------------------------------------------------------===//
// JSON String Processing
//===----------------------------------------------------------------------===//

// Decodes a JSON-escaped string into raw UTF-8 bytes.
static iree_status_t iree_tooling_decode_json_string(
    iree_string_view_t raw_input, iree_allocator_t allocator, char** out_buffer,
    iree_string_view_t* out_text) {
  // Strip surrounding quotes if present.
  iree_string_view_t escaped = raw_input;
  if (escaped.size >= 2 && escaped.data[0] == '"' &&
      escaped.data[escaped.size - 1] == '"') {
    escaped = iree_string_view_substr(escaped, 1, escaped.size - 2);
  }

  // First pass: compute required size.
  iree_host_size_t decoded_length = 0;
  IREE_RETURN_IF_ERROR(
      iree_json_unescape_string(escaped, 0, NULL, &decoded_length));

  // Allocate and decode.
  char* buffer = NULL;
  IREE_RETURN_IF_ERROR(
      iree_allocator_malloc(allocator, decoded_length + 1, (void**)&buffer));
  iree_status_t status = iree_json_unescape_string(escaped, decoded_length + 1,
                                                   buffer, &decoded_length);
  if (!iree_status_is_ok(status)) {
    iree_allocator_free(allocator, buffer);
    return status;
  }
  buffer[decoded_length] = '\0';

  *out_buffer = buffer;
  *out_text = iree_make_string_view(buffer, decoded_length);
  return iree_ok_status();
}

//===----------------------------------------------------------------------===//
// Main
//===----------------------------------------------------------------------===//

int main(int argc, char** argv) {
  IREE_TRACE_APP_ENTER();
  IREE_TRACE_ZONE_BEGIN(z0);

  iree_allocator_t host_allocator = iree_allocator_system();
  int exit_code = EXIT_SUCCESS;

  iree_flags_set_usage(
      "iree-tokenize",
      "Tokenizes text using HuggingFace tokenizer.json files.\n"
      "Outputs comma-separated token IDs (use --json for JSON format).\n"
      "\n"
      "Usage:\n"
      "  iree-tokenize --tokenizer=<file> [flags] <text>\n"
      "\n"
      "Examples:\n"
      "\n"
      "  Encode text to token IDs (default: comma-separated):\n"
      "    iree-tokenize --tokenizer=tokenizer.json \"hello, world!\"\n"
      "    101,7592,1010,2088,999,102\n"
      "\n"
      "  JSON output:\n"
      "    iree-tokenize --tokenizer=tokenizer.json --json \"hello, world!\"\n"
      "    {\"ids\":[101,7592,1010,2088,999,102]}\n"
      "\n"
      "  Use with iree-run-module:\n"
      "    iree-run-module --module=model.vmfb \\\n"
      "      --input=\"6xi32=$(iree-tokenize --tokenizer=tokenizer.json "
      "'hello')\"\n"
      "\n"
      "  Encode without special tokens (no [CLS]/[SEP] or BOS/EOS):\n"
      "    iree-tokenize --tokenizer=tokenizer.json --special=false \"hello "
      "world\"\n"
      "    7592,2088\n"
      "\n"
      "  Show token-to-byte offset mappings:\n"
      "    iree-tokenize --tokenizer=tokenizer.json --offsets \"hello world\"\n"
      "    7592[0:5],2088[6:11]\n"
      "\n"
      "  Decode token IDs back to text:\n"
      "    iree-tokenize --tokenizer=tokenizer.json --decode "
      "\"101,7592,2088,102\"\n"
      "    [CLS]helloworld[SEP]\n"
      "\n"
      "  Show tokenizer info (always JSON):\n"
      "    iree-tokenize --tokenizer=tokenizer.json --info\n"
      "    {\"vocab_size\":30522,\"model_type\":\"BPE\",\"unk_id\":100,"
      "\"cls_id\":101,\"sep_id\":102}\n"
      "\n"
      "  Batch mode - encode one line per input from stdin:\n"
      "    echo -e \"hello\\nworld\" | iree-tokenize "
      "--tokenizer=tokenizer.json --batch\n"
      "    101,7592,102\n"
      "    101,2088,102\n"
      "\n"
      "  Stream mode - continuous stdin encoding (no line buffering):\n"
      "    cat large_file.txt | iree-tokenize --tokenizer=tokenizer.json "
      "--stream\n"
      "    101,7592,...\n"
      "\n"
      "  Truncate output to max length:\n"
      "    iree-tokenize --tokenizer=tokenizer.json --max_length=5 \"hello "
      "world foo\"\n"
      "    101,7592,2088,29379,102\n"
      "\n"
      "  Benchmark encode throughput:\n"
      "    iree-tokenize --tokenizer=tokenizer.json --benchmark=oneshot "
      "\"hello world\"\n"
      "\n"
      "  JSON output with jq:\n"
      "    iree-tokenize --tokenizer=tokenizer.json --json \"hello\" | jq "
      "'.ids'\n"
      "    [101,7592,102]\n");
  iree_flags_parse_checked(IREE_FLAGS_PARSE_MODE_DEFAULT, &argc, &argv);

  if (FLAG_tokenizer[0] == '\0') {
    fprintf(stderr,
            "Error: missing --tokenizer=<file> flag\n"
            "Usage: iree-tokenize --tokenizer=<file> [flags] <text>\n"
            "Run with --help for more information.\n");
    IREE_TRACE_ZONE_END(z0);
    IREE_TRACE_APP_EXIT(EXIT_FAILURE);
    return EXIT_FAILURE;
  }

  iree_string_view_t tokenizer_path = iree_make_cstring_view(FLAG_tokenizer);

  // Load tokenizer file.
  iree_io_file_contents_t* file_contents = NULL;
  iree_status_t status = iree_io_file_contents_map(
      tokenizer_path, IREE_IO_FILE_ACCESS_READ, host_allocator, &file_contents);

  // Create tokenizer from either HuggingFace JSON or tiktoken format.
  iree_tokenizer_t* tokenizer = NULL;
  if (iree_status_is_ok(status)) {
    iree_string_view_t file_data =
        iree_make_string_view((const char*)file_contents->const_buffer.data,
                              file_contents->const_buffer.data_length);

    if (iree_string_view_ends_with(tokenizer_path, IREE_SV(".tiktoken"))) {
      // Tiktoken format: resolve encoding config from --encoding flag or
      // infer from filename (e.g., "cl100k_base.tiktoken" -> "cl100k_base").
      const iree_tokenizer_tiktoken_config_t* config = NULL;
      iree_string_view_t encoding = iree_make_cstring_view(FLAG_encoding);
      if (iree_string_view_is_empty(encoding)) {
        // Extract basename: find last '/' (or '\' on Windows).
        iree_host_size_t last_sep = iree_string_view_find_last_of(
            tokenizer_path, IREE_SV("/\\"), IREE_STRING_VIEW_NPOS);
        encoding = (last_sep != IREE_STRING_VIEW_NPOS)
                       ? iree_string_view_substr(tokenizer_path, last_sep + 1,
                                                 IREE_HOST_SIZE_MAX)
                       : tokenizer_path;
        iree_string_view_consume_suffix(&encoding, IREE_SV(".tiktoken"));
      }
      if (iree_string_view_equal(encoding, IREE_SV("cl100k_base"))) {
        config = iree_tokenizer_tiktoken_config_cl100k_base();
      } else if (iree_string_view_equal(encoding, IREE_SV("o200k_base"))) {
        config = iree_tokenizer_tiktoken_config_o200k_base();
      } else if (iree_string_view_equal(encoding, IREE_SV("r50k_base"))) {
        config = iree_tokenizer_tiktoken_config_r50k_base();
      } else if (iree_string_view_equal(encoding, IREE_SV("p50k_base"))) {
        config = iree_tokenizer_tiktoken_config_p50k_base();
      } else {
        status = iree_make_status(
            IREE_STATUS_INVALID_ARGUMENT,
            "unknown tiktoken encoding '%.*s'; use --encoding with one of: "
            "cl100k_base, o200k_base, r50k_base, p50k_base",
            (int)encoding.size, encoding.data);
      }
      if (iree_status_is_ok(status)) {
        status = iree_tokenizer_from_tiktoken(file_data, config, host_allocator,
                                              &tokenizer);
      }
    } else {
      // Default: HuggingFace JSON format.
      status = iree_tokenizer_from_huggingface_json(file_data, host_allocator,
                                                    &tokenizer);
    }
  }

  // Validate flag combinations.
  if (iree_status_is_ok(status) && FLAG_json_string) {
    if (FLAG_batch || FLAG_stream) {
      fprintf(stderr,
              "Error: --json_string is not supported with --batch/--stream\n"
              "(file/stdin input preserves UTF-8; use --json_string only for "
              "command-line arguments)\n");
      status = iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                                "--json_string requires single input mode");
    } else if (FLAG_decode) {
      fprintf(stderr,
              "Error: --json_string is not supported with --decode\n"
              "(decode takes numeric IDs, not text)\n");
      status = iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                                "--json_string requires encode mode");
    }
  }

  // Process based on mode.
  if (iree_status_is_ok(status)) {
    if (FLAG_info) {
      status = iree_tooling_tokenize_info(tokenizer);
    } else if (FLAG_benchmark[0] != '\0') {
      // Benchmark mode requires text input.
      if (argc < 2) {
        fprintf(stderr, "Error: --benchmark requires input text argument\n");
        status = iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                                  "--benchmark requires input text");
      } else {
        iree_string_view_t input = iree_make_cstring_view(argv[1]);
        char* decoded_buffer = NULL;
        if (FLAG_json_string) {
          status = iree_tooling_decode_json_string(input, host_allocator,
                                                   &decoded_buffer, &input);
        }
        if (iree_status_is_ok(status)) {
          status =
              iree_tooling_tokenize_benchmark(tokenizer, input, host_allocator);
        }
        iree_allocator_free(host_allocator, decoded_buffer);
      }
    } else if (FLAG_stream) {
      if (FLAG_decode) {
        fprintf(stderr, "Error: --stream is not supported with --decode\n");
        status = iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                                  "--stream requires encode mode");
      } else {
        status =
            iree_tooling_tokenize_stdin_streaming(tokenizer, host_allocator);
      }
    } else if (FLAG_batch) {
      status = iree_tooling_tokenize_batch(tokenizer, host_allocator);
    } else if (argc < 2) {
      fprintf(stderr,
              "Error: missing input text\n"
              "Usage: iree-tokenize --tokenizer=<file> [flags] <text>\n");
      status = iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "missing input");
    } else {
      iree_string_view_t input = iree_make_cstring_view(argv[1]);
      char* decoded_buffer = NULL;
      if (FLAG_json_string) {
        status = iree_tooling_decode_json_string(input, host_allocator,
                                                 &decoded_buffer, &input);
      }
      if (iree_status_is_ok(status)) {
        if (FLAG_decode) {
          status =
              iree_tooling_tokenize_decode(tokenizer, input, host_allocator);
        } else {
          status =
              iree_tooling_tokenize_encode(tokenizer, input, host_allocator);
        }
      }
      iree_allocator_free(host_allocator, decoded_buffer);
    }
  }

  // Cleanup.
  if (tokenizer) iree_tokenizer_free(tokenizer);
  if (file_contents) iree_io_file_contents_free(file_contents);

  fflush(stdout);
  if (!iree_status_is_ok(status)) {
    iree_status_fprint(stderr, status);
    iree_status_free(status);
    exit_code = EXIT_FAILURE;
  }
  fflush(stderr);

  IREE_TRACE_ZONE_END(z0);
  IREE_TRACE_APP_EXIT(exit_code);
  return exit_code;
}