// The tokenizer SigLIP 2's text model was trained with: Gemma's, a SentencePiece-style BPE with
// byte fallback, read from its tokenizer.json. What it does, as that file says:
//
//   - spaces become "▁" (the normalizer), and the whole text is then one word;
//   - each character is a piece, or when the vocabulary has no such piece, its UTF-8 bytes as
//     <0xXX> pieces (byte fallback);
//   - neighbouring pieces are merged, the pair with the lowest rank in the merges list first and the
//     leftmost of equal ranks, until no listed pair is left;
//   - <eos> goes at the end, and <pad> fills it out to max_length (64, as SigLIP 2 was trained).
//
// Before all that, text is lowercased, as SigLIP 2 was trained with lowercased text. Only ASCII
// letters are lowered here; other letters keep their case.

Gemma_Tokenizer :: struct {
  text: string;                       // tokenizer.json, which unescaped pieces point into
  vocab: Table(string, s32);          // piece to id
  merges: Table(u64, Merge);          // left id << 32 | right id, to the merge
  byte_ids: [256] s32;                // <0x00>..<0xFF>
  unk, eos, pad: s32;
}

Merge :: struct { rank: s32; result: s32; }

GEMMA_SPACE :: "\xE2\x96\x81";   // ▁, U+2581

// Reads tokenizer.json. Everything from the context allocator; tokenizer_free frees it.
tokenizer_load :: (path: string) -> *Gemma_Tokenizer, ok: bool, error: string {
  text, read := read_entire_file(path, log_errors = false);
  if !read  return null, false, tprint("Cannot read %.", path);
  t := New(Gemma_Tokenizer);
  t.text = text;

  // "vocab": { "<pad>": 0, ... }
  p := find_index_from_left(text, "\"vocab\"");
  if p < 0 { tokenizer_free(t); return null, false, tprint("% has no vocab.", path); }
  rest := advance(text, p + 7);
  if !json_expect(*rest, #char ":") || !json_expect(*rest, #char "{") { tokenizer_free(t); return null, false, tprint("% has a vocab that is not an object.", path); }
  table_ensure_space(*t.vocab, 300_000);
  while true {
    json_skip_space(*rest);
    if rest.count && rest[0] == #char "}" { advance(*rest); break; }
    piece, ok := json_string(*rest);
    if !ok || !json_expect(*rest, #char ":") { tokenizer_free(t); return null, false, tprint("% has a vocab entry that cannot be read.", path); }
    json_skip_space(*rest);
    id, number_ok, after := string_to_int(rest, T = s64);
    if !number_ok { tokenizer_free(t); return null, false, tprint("% has a vocab id that cannot be read.", path); }
    rest = after;
    table_set(*t.vocab, piece, cast(s32) id);
    json_skip_space(*rest);
    if rest.count && rest[0] == #char ","  advance(*rest);
  }

  // "merges": [ ["a", "b"], ... ]
  p = find_index_from_left(rest, "\"merges\"");
  if p < 0 { tokenizer_free(t); return null, false, tprint("% has no merges.", path); }
  rest = advance(rest, p + 8);
  if !json_expect(*rest, #char ":") || !json_expect(*rest, #char "[") { tokenizer_free(t); return null, false, tprint("% has merges that are not a list.", path); }
  table_ensure_space(*t.merges, 800_000);
  rank: s32 = 0;
  joined: [..] u8;   // each merge's two pieces end to end, to look up their result
  defer array_free(joined);
  while true {
    json_skip_space(*rest);
    if rest.count && rest[0] == #char "]" { advance(*rest); break; }
    if !json_expect(*rest, #char "[") { tokenizer_free(t); return null, false, tprint("% has a merge that is not a pair (older \"a b\" merges are not read).", path); }
    left, left_ok := json_string(*rest);
    json_expect(*rest, #char ",");
    right, right_ok := json_string(*rest);
    if !left_ok || !right_ok || !json_expect(*rest, #char "]") { tokenizer_free(t); return null, false, tprint("% has a merge that cannot be read.", path); }
    json_skip_space(*rest);
    if rest.count && rest[0] == #char ","  advance(*rest);

    left_found,   left_id   := table_find(*t.vocab, left);
    right_found,  right_id  := table_find(*t.vocab, right);
    joined.count = 0;
    array_add(*joined, ..cast([] u8) left);
    array_add(*joined, ..cast([] u8) right);
    result_found, result_id := table_find(*t.vocab, cast(string) joined);
    if left_found && right_found && result_found {
      table_set(*t.merges, ((cast(u64) left_id) << 32) | cast(u64) cast(u32) right_id, .{rank, result_id});
    }
    rank += 1;
  }

  found: bool;
  found, t.unk = table_find(*t.vocab, "<unk>");
  if !found { tokenizer_free(t); return null, false, tprint("% has no <unk>.", path); }
  found, t.eos = table_find(*t.vocab, "<eos>");
  if !found { tokenizer_free(t); return null, false, tprint("% has no <eos>.", path); }
  found, t.pad = table_find(*t.vocab, "<pad>");
  if !found { tokenizer_free(t); return null, false, tprint("% has no <pad>.", path); }
  for 0..255 {
    byte_found, id := table_find(*t.vocab, tprint("<0x%>", formatInt(it, base = 16, minimum_digits = 2)));
    t.byte_ids[it] = ifx byte_found then id else t.unk;
  }
  return t, true, "";
}

tokenizer_free :: (t: *Gemma_Tokenizer) {
  if !t  return;
  // Pieces with escapes were copied; the others point into text.
  for _, piece: t.vocab  if !tokenizer_points_into(t, piece)  free(piece);
  table_deinit(*t.vocab);
  table_deinit(*t.merges);
  free(t.text);
  free(t);
}

// text as max_length token ids: its pieces (at most max_length - 1), <eos>, then <pad>. Temp.
tokenizer_encode :: (t: *Gemma_Tokenizer, text: string, max_length := 64) -> [] s64 {
  // Lowercased, spaces as ▁.
  normal: String_Builder;
  normal.allocator = temp;
  for 0..text.count - 1 {
    c := text[it];
    if c == #char " "                          append(*normal, GEMMA_SPACE);
    else if c >= #char "A" && c <= #char "Z"   append(*normal, c + (#char "a" - #char "A"));
    else                                       append(*normal, c);
  }
  word := builder_to_string(*normal,, temp);

  // One piece per character; byte pieces for those the vocabulary lacks.
  symbols: [..] s32;
  symbols.allocator = temp;
  i := 0;
  while i < word.count {
    length := utf8_length(word[i]);
    if i + length > word.count  length = word.count - i;   // a cut-off sequence: its bytes
    found, id := table_find(*t.vocab, slice(word, i, length));
    if found  array_add(*symbols, id);
    else      for b: 0..length - 1  array_add(*symbols, t.byte_ids[word[i + b]]);
    i += length;
  }

  // Merges, lowest rank first, the leftmost of equal ranks.
  while symbols.count > 1 {
    best := -1;
    best_merge: Merge;
    for 0..symbols.count - 2 {
      found, merge := table_find(*t.merges, ((cast(u64) symbols[it]) << 32) | cast(u64) cast(u32) symbols[it + 1]);
      if found && (best < 0 || merge.rank < best_merge.rank) { best = it; best_merge = merge; }
    }
    if best < 0  break;
    symbols[best] = best_merge.result;
    array_ordered_remove_by_index(*symbols, best + 1);
  }

  ids := NewArray(max_length, s64,, temp);
  kept := min(symbols.count, max_length - 1);
  for 0..kept - 1  ids[it] = symbols[it];
  ids[kept] = t.eos;
  for kept + 1..max_length - 1  ids[it] = t.pad;
  return ids;
}

#scope_file

// The length of the UTF-8 sequence a byte starts: 1 for ASCII and stray continuation bytes.
utf8_length :: (lead: u8) -> s64 {
  if lead < 0xC0  return 1;
  if lead < 0xE0  return 2;
  if lead < 0xF0  return 3;
  return 4;
}

tokenizer_points_into :: (t: *Gemma_Tokenizer, s: string) -> bool {
  return s.data >= t.text.data && s.data < t.text.data + t.text.count;
}

json_skip_space :: (s: *string) {
  while s.count && (s.*[0] == #char " " || s.*[0] == #char "\n" || s.*[0] == #char "\r" || s.*[0] == #char "\t")  advance(s);
}

json_expect :: (s: *string, c: u8) -> bool {
  json_skip_space(s);
  if !s.count || s.*[0] != c  return false;
  advance(s);
  return true;
}

// A JSON string at the start of s: a view into s when it has no escapes, else unescaped into a
// copy from the context allocator.
json_string :: (s: *string) -> string, ok: bool {
  json_skip_space(s);
  if !s.count || s.*[0] != #char "\""  return "", false;
  advance(s);
  end := 0;
  escaped := false;
  while end < s.count && s.*[end] != #char "\"" {
    if s.*[end] == #char "\\" { escaped = true; end += 1; }
    end += 1;
  }
  if end >= s.count  return "", false;
  raw := slice(s.*, 0, end);
  advance(s, end + 1);
  if !escaped  return raw, true;

  builder: String_Builder;
  i := 0;
  while i < raw.count {
    c := raw[i];
    if c != #char "\\" { append(*builder, c); i += 1; continue; }
    if i + 1 >= raw.count  return "", false;
    e := raw[i + 1];
    i += 2;
    if e == {
      case #char "n";  append(*builder, #char "\n");
      case #char "t";  append(*builder, #char "\t");
      case #char "r";  append(*builder, #char "\r");
      case #char "b";  append(*builder, 8);
      case #char "f";  append(*builder, 12);
      case #char "u";
        code, ok := hex4(raw, i);
        if !ok  return "", false;
        i += 4;
        if code >= 0xD800 && code < 0xDC00 && i + 6 <= raw.count && raw[i] == #char "\\" && raw[i + 1] == #char "u" {
          low, low_ok := hex4(raw, i + 2);
          if low_ok && low >= 0xDC00 && low < 0xE000 { code = 0x10000 + ((code - 0xD800) << 10) + (low - 0xDC00); i += 6; }
        }
        append_utf8(*builder, code);
      case;  append(*builder, e);   // \" \\ \/
    }
  }
  return builder_to_string(*builder), true;
}

hex4 :: (s: string, at: s64) -> u32, ok: bool {
  if at + 4 > s.count  return 0, false;
  value: u32;
  for 0..3 {
    c := s[at + it];
    digit: u32;
    if c >= #char "0" && c <= #char "9"       digit = c - #char "0";
    else if c >= #char "a" && c <= #char "f"  digit = c - #char "a" + 10;
    else if c >= #char "A" && c <= #char "F"  digit = c - #char "A" + 10;
    else return 0, false;
    value = value * 16 + digit;
  }
  return value, true;
}

append_utf8 :: (builder: *String_Builder, code: u32) {
  if code < 0x80         { append(builder, cast(u8) code); return; }
  if code < 0x800        { append(builder, cast(u8) (0xC0 | (code >> 6))); append(builder, cast(u8) (0x80 | (code & 0x3F))); return; }
  if code < 0x10000      { append(builder, cast(u8) (0xE0 | (code >> 12))); append(builder, cast(u8) (0x80 | ((code >> 6) & 0x3F))); append(builder, cast(u8) (0x80 | (code & 0x3F))); return; }
  append(builder, cast(u8) (0xF0 | (code >> 18))); append(builder, cast(u8) (0x80 | ((code >> 12) & 0x3F)));
  append(builder, cast(u8) (0x80 | ((code >> 6) & 0x3F))); append(builder, cast(u8) (0x80 | (code & 0x3F)));
}