/* 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>
#include <pthread.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 pthread_once_t tok_once = PTHREAD_ONCE_INIT;
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) {
  unsigned int size = 1;
  while (size < (unsigned int)tok_ntok * 2u) size <<= 1;
  const char **keys = (const char **)calloc(size, sizeof(char *));
  int *ranks = (int *)calloc(size, sizeof(int));
  int *klen = (int *)calloc(size, sizeof(int));
  if (!keys || !ranks || !klen) {
    free(keys);
    free(ranks);
    free(klen);
    return;
  }
  unsigned int mask = size - 1;
  int offset = 0;
  for (int i = 0; i < tok_ntok; i++) {
    int kl = tok_len[i];
    const char *key = (const char *)(tok_blob + offset);
    unsigned long h = tok_fnv1a((const unsigned char *)key, kl) & mask;
    while (keys[h]) {
      h = (h + 1) & mask;
    }
    keys[h] = key;
    ranks[h] = tok_rank[i];
    klen[h] = kl;
    offset += kl;
  }
  tok_mask = mask;
  tok_ranks = ranks;
  tok_klen = klen;
  tok_keys = keys;
}

static int tok_lookup(const unsigned char *s, int len) {
  pthread_once(&tok_once, 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 <= 0x02af) {
    return (cp != 0x00d7 && cp != 0x00f7);        /* Latin-1 + Ext-A/B + IPA (excl. × ÷) */
  }
  if (cp >= 0x0370 && cp <= 0x03ff) {
    return (cp != 0x037e && cp != 0x0387);        /* Greek */
  }
  if (cp >= 0x0400 && cp <= 0x052f) return 1;     /* Cyrillic + Supplement */
  if (cp >= 0x0531 && cp <= 0x0587) return 1;     /* Armenian */
  if (cp >= 0x05d0 && cp <= 0x05ea) return 1;     /* Hebrew */
  if ((cp >= 0x0620 && cp <= 0x064a) ||
      (cp >= 0x066e && cp <= 0x06d3) ||
      (cp >= 0x06ee && cp <= 0x06ff) ||
      (cp >= 0x0750 && cp <= 0x077f) ||
      (cp >= 0x08a0 && cp <= 0x08ff)) return 1;   /* Arabic */
  if (cp >= 0x0904 && cp <= 0x0d7f) {
    return !(cp >= 0x0964 && cp <= 0x096f);       /* Devanagari + Indic scripts */
  }
  if ((cp >= 0x0e01 && cp <= 0x0e30) ||
      (cp >= 0x0e32 && cp <= 0x0e33) ||
      (cp >= 0x0e40 && cp <= 0x0e46)) return 1;   /* Thai */
  if (cp >= 0x10a0 && cp <= 0x10ff) return 1;     /* Georgian */
  if (cp >= 0x1100 && cp <= 0x11ff) return 1;     /* Hangul Jamo */
  if (cp >= 0x1e00 && cp <= 0x1fff) return 1;     /* Latin Ext Additional + Greek Ext */
  if (cp >= 0x3041 && cp <= 0x3096) return 1;     /* Hiragana */
  if (cp >= 0x30a1 && cp <= 0x30fc && cp != 0x30fb) return 1; /* Katakana (excl. ・) */
  if (cp >= 0x3131 && cp <= 0x318e) return 1;     /* Hangul Compatibility Jamo */
  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 >= 0xf900 && cp <= 0xfaff) return 1;     /* CJK Compatibility Ideographs */
  if (cp >= 0xff21 && cp <= 0xff3a) return 1;     /* fullwidth A-Z */
  if (cp >= 0xff41 && cp <= 0xff5a) return 1;     /* fullwidth a-z */
  if (cp >= 0x20000 && cp <= 0x2fa1f) return 1;   /* CJK Ext B-F + Compat */
  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 || (t[pos + j] & 0xc0) != 0x80) { 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 = (!isS && !isL && !isN);
  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: \s++$ | \s*[\r\n] | \s+(?!\S) | \s */
  int e = i;
  int last_ws_len = cplen;
  int last_nl_end = 0;
  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) last_nl_end = e + cl;
    e += cl;
  }
  if (e == len) return e;                        /* 5) \s++$ (trailing whitespace) */
  if (last_nl_end > i) return last_nl_end;       /* 6) \s*[\r\n] (up to last CR/LF) */
  if ((e - i) > last_ws_len) return e - last_ws_len; /* 7) \s+(?!\S) (leave last ws) */
  return e;                                      /* 8) \s (single whitespace) */
}

/* ----------------------------- BPE ------------------------------------ */
#define MAX_TOKEN_BYTES 128

typedef struct {
  int prev;
  int next;
  int rank;
} bpe_node_t;

typedef struct {
  int rank;
  int pos;
} bpe_heap_item_t;

static int heap_less(bpe_heap_item_t a, bpe_heap_item_t b) {
  return (a.rank < b.rank) || (a.rank == b.rank && a.pos < b.pos);
}

static void heap_push(bpe_heap_item_t *h, int *hn, bpe_heap_item_t v) {
  int i = (*hn)++;
  while (i > 0) {
    int p = (i - 1) >> 1;
    if (!heap_less(v, h[p])) break;
    h[i] = h[p];
    i = p;
  }
  h[i] = v;
}

static bpe_heap_item_t heap_pop(bpe_heap_item_t *h, int *hn) {
  bpe_heap_item_t top = h[0];
  bpe_heap_item_t last = h[--(*hn)];
  int n = *hn, i = 0;
  while ((i << 1) + 1 < n) {
    int l = (i << 1) + 1, r = l + 1;
    int c = (r < n && heap_less(h[r], h[l])) ? r : l;
    if (!heap_less(h[c], last)) break;
    h[i] = h[c];
    i = c;
  }
  if (n > 0) h[i] = last;
  return top;
}

/* Count tokens in a single raw-byte piece via O(N log N) lowest-rank merges. */
static int tok_bpe(const unsigned char *piece, int plen) {
  if (plen <= 0) return 0;
  if (plen == 1) return 1;
  bpe_node_t *nodes = (bpe_node_t *)malloc(sizeof(bpe_node_t) * (size_t)plen);
  bpe_heap_item_t *heap = (bpe_heap_item_t *)malloc(sizeof(bpe_heap_item_t) * (size_t)plen * 2u);
  if (!nodes || !heap) {
    free(nodes);
    free(heap);
    return (plen + 3) / 4;
  }
  int hn = 0;
  for (int k = 0; k < plen; k++) {
    nodes[k].prev = k - 1;
    nodes[k].next = k + 1;
    if (k + 1 < plen) {
      int r = tok_lookup(piece + k, 2);
      nodes[k].rank = r;
      if (r >= 0) {
        bpe_heap_item_t item = { r, k };
        heap_push(heap, &hn, item);
      }
    } else {
      nodes[k].rank = -1;
    }
  }
  int nsym = plen;
  while (nsym > 1 && hn > 0) {
    bpe_heap_item_t top = heap_pop(heap, &hn);
    int best = top.pos;
    int right = nodes[best].next;
    if (right >= plen || nodes[best].rank != top.rank) continue;

    int nr = nodes[right].next;
    nodes[best].next = nr;
    if (nr < plen) nodes[nr].prev = best;
    nodes[right].next = plen + 1;
    nodes[right].rank = -1;
    nsym--;

    if (nr < plen) {
      int clen = nodes[nr].next - best;
      int r = (clen <= MAX_TOKEN_BYTES) ? tok_lookup(piece + best, clen) : -1;
      nodes[best].rank = r;
      if (r >= 0) {
        bpe_heap_item_t item = { r, best };
        heap_push(heap, &hn, item);
      }
    } else {
      nodes[best].rank = -1;
    }

    int left = nodes[best].prev;
    if (left >= 0) {
      int clen = nr - left;
      int r = (clen <= MAX_TOKEN_BYTES) ? tok_lookup(piece + left, clen) : -1;
      nodes[left].rank = r;
      if (r >= 0) {
        bpe_heap_item_t item = { r, left };
        heap_push(heap, &hn, item);
      }
    }
  }
  free(nodes);
  free(heap);
  return nsym;
}

/* ----------------------------- public --------------------------------- */
int tok_count(const char *text, int len) {
  if (!text || len <= 0) return 0;
  pthread_once(&tok_once, 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);
    total += tok_bpe(t + i, end - i);
    i = end;
  }
  return total;
}

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