/* tokenizer.c - embedded cl100k_base BPE tokenizer (pure C, no deps).
 *
 * Counts tokens for arbitrary UTF-8 text using the OpenAI cl100k_base
 * vocabulary (embedded in tokenizer_data.c). The model served by this
 * gateway (Hunyuan) uses a different tokenizer, but cl100k_base is the
 * canonical multilingual OpenAI tokenizer and gives deterministic,
 * real subword counts that are far closer to the truth than a naive
 * bytes/4 heuristic.
 *
 * The pre-tokenizer mirrors tiktoken's cl100k regex:
 *   '(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}++|\p{N}{1,3}+|
 *     ?[^\s\p{L}\p{N}]++[\r\n]*+|\s++$|\s*[\r\n]|\s+(?!\S)|\s
 * Unicode letter/number classification is approximated with compact
 * codepoint-range tables (covers Latin, CJK, Kana, Hangul, Cyrillic,
 * Greek and common digits); this only affects script boundaries, not
 * the BPE merge counts for in-script text.
 */
#include "hygate.h"
#include <string.h>
#include <stdint.h>
#include <stdlib.h>

extern const int tok_ntok;
extern const unsigned char tok_blob[];
extern const int tok_rank[];
extern const uint16_t tok_len[];

/* ----------------------------- vocab hash ----------------------------- */
static int tok_inited = 0;
static const char **tok_keys = NULL;
static int *tok_ranks = NULL;
static int *tok_klen = NULL;static unsigned int tok_mask = 0;

static unsigned long tok_fnv1a(const unsigned char *s, int len) {
  unsigned long h = 1469598103934665603UL;
  for (int i = 0; i < len; i++) {
    h ^= (unsigned long)s[i];
    h *= 1099511628211UL;
  }
  return h;
}

static void tok_init(void) {
  if (tok_inited) return;
  tok_inited = 1;
  unsigned int size = 1;
  while (size < (unsigned int)tok_ntok * 2u) size <<= 1;
  tok_mask = size - 1;
  tok_keys = (const char **)calloc(size, sizeof(char *));
  tok_ranks = (int *)calloc(size, sizeof(int));
  tok_klen = (int *)calloc(size, sizeof(int));
  if (!tok_keys || !tok_ranks || !tok_klen) return;
  int offset = 0;
  for (int i = 0; i < tok_ntok; i++) {
    int klen = tok_len[i];
    const char *key = (const char *)(tok_blob + offset);
    unsigned long h = tok_fnv1a((const unsigned char *)key, klen) & tok_mask;
    while (tok_keys[h]) {
      const char *k = tok_keys[h];
      if (tok_klen[h] == klen && memcmp(k, key, klen) == 0) break; /* dup (none expected) */
      h = (h + 1) & tok_mask;
    }
    tok_keys[h] = key;
    tok_ranks[h] = tok_rank[i];
    tok_klen[h] = klen;
    offset += klen;
  }
}

static int tok_lookup(const unsigned char *s, int len) {
  if (!tok_inited) tok_init();
  if (!tok_keys) return -1;
  unsigned long h = tok_fnv1a(s, len) & tok_mask;
  while (tok_keys[h]) {
    const char *k = tok_keys[h];
    if (tok_klen[h] == len && memcmp(k, s, len) == 0) return tok_ranks[h];
    h = (h + 1) & tok_mask;
  }
  return -1;
}

/* ----------------------------- unicode -------------------------------- */
static int u_is_space(int cp) {
  switch (cp) {
    case 0x20: case 0x09: case 0x0a: case 0x0d: case 0x0b: case 0x0c:
    case 0xa0: case 0x1680: case 0x2028: case 0x2029: case 0x202f:
    case 0x205f: case 0x3000:
      return 1;
  }
  if (cp >= 0x2000 && cp <= 0x200a) return 1;
  return 0;
}
static int u_is_number(int cp) {
  if (cp >= 0x30 && cp <= 0x39) return 1;        /* ASCII digits */
  if (cp >= 0xff10 && cp <= 0xff19) return 1;     /* fullwidth digits */
  if (cp >= 0x0660 && cp <= 0x0669) return 1;     /* Arabic-Indic */
  if (cp >= 0x06f0 && cp <= 0x06f9) return 1;     /* Extended Arabic-Indic */
  if (cp >= 0x0966 && cp <= 0x096f) return 1;     /* Devanagari */
  return 0;
}
static int u_is_letter(int cp) {
  if ((cp >= 'a' && cp <= 'z') || (cp >= 'A' && cp <= 'Z')) return 1;
  if (cp >= 0x00c0 && cp <= 0x024f) return 1;     /* Latin-1 + Latin Extended-A/B */
  if (cp >= 0x0370 && cp <= 0x03ff) return 1;     /* Greek */
  if (cp >= 0x0400 && cp <= 0x04ff) return 1;     /* Cyrillic */
  if (cp >= 0x3040 && cp <= 0x30ff) return 1;     /* Hiragana + Katakana */
  if (cp >= 0x3400 && cp <= 0x4dbf) return 1;     /* CJK Ext A */
  if (cp >= 0x4e00 && cp <= 0x9fff) return 1;     /* CJK Unified Ideographs */
  if (cp >= 0xac00 && cp <= 0xd7a3) return 1;     /* Hangul Syllables */
  if (cp >= 0xff21 && cp <= 0xff5a) return 1;     /* fullwidth A-Z */
  if (cp >= 0xff41 && cp <= 0xff5a) return 1;     /* fullwidth a-z */
  return 0;
}

/* Decode one UTF-8 codepoint at t[pos]; sets *cp and *cplen. */
static void decode_cp(const unsigned char *t, int len, int pos, int *cp, int *cplen) {
  unsigned char b = t[pos];
  int c, l;
  if (b < 0x80) { c = b; l = 1; }
  else if ((b >> 5) == 0x6) { c = b & 0x1f; l = 2; }
  else if ((b >> 4) == 0xe) { c = b & 0x0f; l = 3; }
  else if ((b >> 3) == 0x1e) { c = b & 0x07; l = 4; }
  else { c = b; l = 1; } /* invalid lead byte: treat as single byte */
  for (int j = 1; j < l; j++) {
    if (pos + j >= len) { c = b; l = 1; break; }
    c = (c << 6) | (t[pos + j] & 0x3f);
  }
  *cp = c; *cplen = l;
}

/* ----------------------------- pre-tokenizer -------------------------- */
/* Returns the exclusive end byte offset of the token piece that starts at
 * position i, implementing the cl100k regex alternatives in order. */
static int tok_piece_end(const unsigned char *t, int len, int i) {
  int cp, cplen;
  decode_cp(t, len, i, &cp, &cplen);
  int ni = i + cplen;
  int ncp = 0, ncl = 0;
  if (ni < len) decode_cp(t, len, ni, &ncp, &ncl);

  /* 1) '(?i:[sdmt]|ll|ve|re) */
  if (cp == 0x27) {
    if (ni < len) {
      int j2 = ni + ncl;
      if (j2 < len) {
        int c2, l2; decode_cp(t, len, j2, &c2, &l2);
        int a = ncp, b = c2;
        if (a >= 'A' && a <= 'Z') a += 32;
        if (b >= 'A' && b <= 'Z') b += 32;
        if ((a == 'l' && b == 'l') || (a == 'v' && b == 'e') || (a == 'r' && b == 'e'))
          return j2 + l2;
      }
      int a = ncp; if (a >= 'A' && a <= 'Z') a += 32;
      if (a == 's' || a == 'd' || a == 'm' || a == 't') return ni + ncl;
    }
    /* apostrophe alone falls through to whitespace/other handling */
  }

  int isL = u_is_letter(cp), isN = u_is_number(cp), isS = u_is_space(cp);
  int is_not_cln = (!isL && !isN && cp != 0x0d && cp != 0x0a);

  /* 2) [^\r\n\p{L}\p{N}]?+\p{L}++ */
  if (isL) {
    int e = i;
    while (e < len) { int c, cl; decode_cp(t, len, e, &c, &cl); if (!u_is_letter(c)) break; e += cl; }
    return e;
  }
  if (is_not_cln && ni < len && u_is_letter(ncp)) {
    int e = ni;
    while (e < len) { int c, cl; decode_cp(t, len, e, &c, &cl); if (!u_is_letter(c)) break; e += cl; }
    return e;
  }

  /* 3) \p{N}{1,3}+ */
  if (isN) {
    int e = i, cnt = 0;
    while (e < len && cnt < 3) {
      int c, cl; decode_cp(t, len, e, &c, &cl);
      if (!u_is_number(c)) break;
      e += cl; cnt++;
    }
    return e;
  }

  /* 4)  ?[^\s\p{L}\p{N}]++[\r\n]*+  (leading space optional) */
  int is_other = (!u_is_space(cp) && !u_is_letter(cp) && !u_is_number(cp));
  if (cp == 0x20 && ni < len && (!u_is_space(ncp) && !u_is_letter(ncp) && !u_is_number(ncp))) {
    int e = ni;
    while (e < len) {
      int c, cl; decode_cp(t, len, e, &c, &cl);
      if (u_is_space(c) || u_is_letter(c) || u_is_number(c)) break;
      e += cl;
    }
    while (e < len) {
      int c, cl; decode_cp(t, len, e, &c, &cl);
      if (c != 0x0d && c != 0x0a) break;
      e += cl;
    }
    return e;
  }
  if (is_other) {
    int e = i;
    while (e < len) {
      int c, cl; decode_cp(t, len, e, &c, &cl);
      if (u_is_space(c) || u_is_letter(c) || u_is_number(c)) break;
      e += cl;
    }
    while (e < len) {
      int c, cl; decode_cp(t, len, e, &c, &cl);
      if (c != 0x0d && c != 0x0a) break;
      e += cl;
    }
    return e;
  }

  /* 5-8) whitespace runs */
  if (isS) {
    int e = i;
    int last_ws_len = cplen;
    int has_nl = (cp == 0x0d || cp == 0x0a);
    while (e < len) {
      int c, cl; decode_cp(t, len, e, &c, &cl);
      if (!u_is_space(c)) break;
      last_ws_len = cl;
      if (c == 0x0d || c == 0x0a) has_nl = 1;
      e += cl;
    }
    if (e == len) return e;                 /* 5) trailing whitespace */
    if (has_nl) return e;                   /* 6) whitespace containing a newline */
    if ((e - i) > last_ws_len) return e - last_ws_len; /* 7) leave last whitespace */
    return e;                               /* 8) single whitespace */
  }

  return i + cplen; /* fallback */
}

/* ----------------------------- BPE ------------------------------------ */
/* Count tokens in a single raw-byte piece via greedy lowest-rank merges. */
static int tok_bpe(const unsigned char *piece, int plen) {
  if (plen <= 0) return 0;
  if (plen > 4096) return plen / 4; /* pathological guard; estimate */
  char **sym = (char **)malloc(sizeof(char *) * (size_t)plen);
  if (!sym) return plen / 4;
  int nsym = 0;
  for (int k = 0; k < plen; k++) {
    char *s = (char *)malloc(2);
    if (!s) { for (int j = 0; j < nsym; j++) free(sym[j]); free(sym); return plen / 4; }
    s[0] = (char)piece[k]; s[1] = 0;
    sym[nsym++] = s;
  }
  char buf[256];
  while (nsym > 1) {
    int best = -1, bestr = 0x7fffffff;
    for (int k = 0; k < nsym - 1; k++) {
      int l1 = (int)strlen(sym[k]), l2 = (int)strlen(sym[k + 1]);
      if (l1 + l2 >= (int)sizeof(buf)) continue;
      memcpy(buf, sym[k], l1);
      memcpy(buf + l1, sym[k + 1], l2);
      buf[l1 + l2] = 0;
      int r = tok_lookup((const unsigned char *)buf, l1 + l2);
      if (r >= 0 && r < bestr) { bestr = r; best = k; }
    }
    if (best < 0) break;
    int l1 = (int)strlen(sym[best]), l2 = (int)strlen(sym[best + 1]);
    char *m = (char *)malloc((size_t)(l1 + l2 + 1));
    memcpy(m, sym[best], l1);
    memcpy(m + l1, sym[best + 1], l2);
    m[l1 + l2] = 0;
    free(sym[best]);
    free(sym[best + 1]);
    sym[best] = m;
    for (int k = best + 1; k < nsym - 1; k++) sym[k] = sym[k + 1];
    nsym--;
  }
  int n = nsym;
  for (int k = 0; k < nsym; k++) free(sym[k]);
  free(sym);
  return n;
}

/* ----------------------------- public --------------------------------- */
int tok_count(const char *text, int len) {
  if (!text || len <= 0) return 0;
  if (!tok_inited) tok_init();
  const unsigned char *t = (const unsigned char *)text;
  int total = 0, i = 0;
  while (i < len) {
    int end = tok_piece_end(t, len, i);
    if (end <= i) end = i + 1;
    total += tok_bpe(t + i, end - i);
    i = end;
  }
  return total;
}

int tok_count_str(const char *text) {
  return tok_count(text, (int)strlen(text));
}