/*
 * §Fase 25.g — BPE merge engine implementation.
 *
 * See bpe.h for the full pillar / mathematics rationale.
 *
 * This file is organised in five sections:
 *
 *   1. Embedded default blobs (#embed primary path + NULL fallback)
 *   2. Hash table (open-addressed FNV-1a, byte-key → rank)
 *   3. Encoder lifecycle (load / destroy / lookups)
 *   4. byte_pair_merge — the canonical tiktoken algorithm
 *   5. UTF-8 boundary + SIMD codepoint counter
 */

#include "bpe.h"

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

/* ──────────────────────────────────────────────────────────────────────
 * 1. Embedded default blobs
 *
 * Primary: C23 `#embed`. Available when the host toolchain has shipped
 * the directive (gcc ≥15, clang ≥19; MSVC has not as of 2026-05-08).
 * Selected with the `__has_embed` predicate that the standard requires
 * compilers to provide whenever `#embed` is implemented.
 *
 * Fallback: `extern` symbols populated by `merges_tables.c` (auto-
 * generated by `build.rs` from the .bin files). The xxd-style array
 * is byte-equivalent to the embed result. Both paths produce the
 * same encoder behaviour — the drift gate verifies it.
 * ────────────────────────────────────────────────────────────────── */

/* `AXON_CSYS_BPE_HAS_C23_EMBED` is defined by `build.rs` after probing
 * the host compiler. Using a build-time probe (rather than the C-side
 * `__has_embed` predicate) sidesteps preprocessor parse-error risk on
 * compilers that haven't shipped `__has_embed` recognition yet. */
#ifdef AXON_CSYS_BPE_HAS_C23_EMBED
static const unsigned char AXON_CSYS_BPE_CL100K_EMBED[] = {
#  embed "merges_cl100k_base.bin"
};
static const unsigned char AXON_CSYS_BPE_O200K_EMBED[] = {
#  embed "merges_o200k_base.bin"
};

const uint8_t* axon_csys_bpe_embedded_cl100k_base(size_t* out_len) {
    if (out_len != NULL) {
        *out_len = sizeof AXON_CSYS_BPE_CL100K_EMBED;
    }
    return (const uint8_t*) AXON_CSYS_BPE_CL100K_EMBED;
}

const uint8_t* axon_csys_bpe_embedded_o200k_base(size_t* out_len) {
    if (out_len != NULL) {
        *out_len = sizeof AXON_CSYS_BPE_O200K_EMBED;
    }
    return (const uint8_t*) AXON_CSYS_BPE_O200K_EMBED;
}

bool axon_csys_bpe_used_c23_embed(void) { return true; }
#else
/* Fallback path — the C kernel does not embed the merges blobs at
 * compile time on this toolchain. The Rust shim provides the bytes
 * via `include_bytes!` and feeds them to `axon_csys_bpe_load`. The
 * `embedded_*` accessors below report (NULL, 0) so callers can detect
 * the path and route accordingly. This is a CI-time signal, not a
 * runtime branch in the encode hot loop. */
const uint8_t* axon_csys_bpe_embedded_cl100k_base(size_t* out_len) {
    if (out_len != NULL) {
        *out_len = 0u;
    }
    return NULL;
}

const uint8_t* axon_csys_bpe_embedded_o200k_base(size_t* out_len) {
    if (out_len != NULL) {
        *out_len = 0u;
    }
    return NULL;
}

bool axon_csys_bpe_used_c23_embed(void) { return false; }
#endif

/* ──────────────────────────────────────────────────────────────────────
 * 2. Hash table — open-addressed FNV-1a, (bytes → blob_offset)
 *
 * Each bucket stores a u32 offset into the source blob. The offset
 * points at the entry's length byte (so the actual bytes are at
 * offset+1, the rank at offset+1+len). UINT32_MAX = empty bucket.
 *
 * Capacity is next_pow2(vocab_size * 2) — load factor ≤0.5 keeps
 * probe lengths short. For cl100k (100256 entries) capacity is
 * 262144; for o200k (199998 entries) it's 524288. Memory cost
 * 1MB and 2MB respectively — acceptable for a one-time-loaded
 * static encoder.
 *
 * FNV-1a chosen for: small constant footprint (~10 LOC), no
 * dependencies, well-distributed for short keys. Not cryptographic
 * — collisions in a public hash are not a security concern here
 * because the keyspace is fixed at table-load time.
 * ────────────────────────────────────────────────────────────────── */

#define AXON_CSYS_BPE_FNV_OFFSET_BASIS 0xcbf29ce484222325ULL
#define AXON_CSYS_BPE_FNV_PRIME        0x100000001b3ULL
#define AXON_CSYS_BPE_HASH_EMPTY       UINT32_MAX

static uint64_t axon_csys_bpe_fnv1a(const uint8_t* bytes, size_t len) {
    uint64_t h = AXON_CSYS_BPE_FNV_OFFSET_BASIS;
    for (size_t i = 0; i < len; ++i) {
        h ^= (uint64_t) bytes[i];
        h *= AXON_CSYS_BPE_FNV_PRIME;
    }
    return h;
}

static uint32_t axon_csys_bpe_next_pow2_u32(uint32_t v) {
    /* Lift v to the smallest power of 2 that is ≥ v and ≥ 16 (no
     * pathological tiny tables). Standard bit-twiddle. */
    if (v < 16u) {
        return 16u;
    }
    --v;
    v |= v >> 1;
    v |= v >> 2;
    v |= v >> 4;
    v |= v >> 8;
    v |= v >> 16;
    return v + 1u;
}

/* ──────────────────────────────────────────────────────────────────────
 * Encoder struct
 * ────────────────────────────────────────────────────────────────── */

struct AxonCsysBpeEncoder {
    const uint8_t* blob;
    size_t blob_len;

    uint32_t vocab_size;

    /* Hash table: bytes → blob_offset (offset to length byte). Each
     * entry is a u32; empty buckets carry UINT32_MAX. */
    uint32_t* hash;
    uint32_t hash_cap;     /* power of 2 */
    uint32_t hash_mask;    /* hash_cap - 1 */

    /* Reverse table: rank → blob_offset. Length = vocab_size. */
    uint32_t* rank_to_offset;

    /* Pretokeniser regex pattern (UTF-8, borrowed from blob). */
    const uint8_t* regex_pat;
    uint32_t regex_pat_len;
};

/* Read little-endian u32 at `p` without alignment assumptions. */
static uint32_t axon_csys_bpe_read_u32_le(const uint8_t* p) {
    return ((uint32_t) p[0])
         | ((uint32_t) p[1] << 8)
         | ((uint32_t) p[2] << 16)
         | ((uint32_t) p[3] << 24);
}

/* Compare bytes at offset `a` (in blob) with `b` (caller-supplied). */
static bool axon_csys_bpe_entry_matches(
    const uint8_t* blob,
    uint32_t entry_off,
    const uint8_t* needle,
    size_t needle_len)
{
    uint8_t entry_len = blob[entry_off];
    if ((size_t) entry_len != needle_len) {
        return false;
    }
    return memcmp(blob + entry_off + 1u, needle, needle_len) == 0;
}

/* Probe the open-addressed table for `bytes[0..len]`.
 * Returns AXON_CSYS_BPE_HASH_EMPTY if absent, otherwise the bucket
 * index containing the matching entry. */
static uint32_t axon_csys_bpe_hash_probe(
    const AxonCsysBpeEncoder* enc,
    const uint8_t* bytes,
    size_t len)
{
    uint64_t h = axon_csys_bpe_fnv1a(bytes, len);
    uint32_t mask = enc->hash_mask;
    uint32_t idx = (uint32_t) h & mask;
    /* Linear probing — bounded by hash_cap iterations to detect a
     * fully-saturated table (cannot happen here because load ≤0.5). */
    for (uint32_t step = 0; step < enc->hash_cap; ++step) {
        uint32_t bucket = enc->hash[idx];
        if (bucket == AXON_CSYS_BPE_HASH_EMPTY) {
            return AXON_CSYS_BPE_HASH_EMPTY;
        }
        if (axon_csys_bpe_entry_matches(enc->blob, bucket, bytes, len)) {
            return idx;
        }
        idx = (idx + 1u) & mask;
    }
    return AXON_CSYS_BPE_HASH_EMPTY;
}

/* Insert at table-load time. Caller guarantees the key is unique. */
static void axon_csys_bpe_hash_insert_unchecked(
    AxonCsysBpeEncoder* enc,
    uint32_t entry_off)
{
    uint8_t key_len = enc->blob[entry_off];
    const uint8_t* key = enc->blob + entry_off + 1u;
    uint64_t h = axon_csys_bpe_fnv1a(key, (size_t) key_len);
    uint32_t mask = enc->hash_mask;
    uint32_t idx = (uint32_t) h & mask;
    while (enc->hash[idx] != AXON_CSYS_BPE_HASH_EMPTY) {
        idx = (idx + 1u) & mask;
    }
    enc->hash[idx] = entry_off;
}

/* ──────────────────────────────────────────────────────────────────────
 * 3. Encoder lifecycle
 * ────────────────────────────────────────────────────────────────── */

AxonCsysBpeEncoder* axon_csys_bpe_load(
    const uint8_t* blob,
    size_t blob_len,
    AxonCsysBpeError* out_err)
{
#define FAIL(code) do { if (out_err) { *out_err = (code); } return NULL; } while (0)

    if (blob == NULL || blob_len < 20u) {
        FAIL(AXON_CSYS_BPE_NULL_ARG);
    }

    /* Header: magic ("AXBP"), version, vocab_size, regex_pat_len. */
    if (!(blob[0] == 'A' && blob[1] == 'X' && blob[2] == 'B' && blob[3] == 'P')) {
        FAIL(AXON_CSYS_BPE_BAD_MAGIC);
    }
    uint32_t version = axon_csys_bpe_read_u32_le(blob + 4u);
    if (version != 1u) {
        FAIL(AXON_CSYS_BPE_BAD_VERSION);
    }
    uint32_t vocab_size = axon_csys_bpe_read_u32_le(blob + 8u);
    uint32_t pat_len = axon_csys_bpe_read_u32_le(blob + 12u);
    if (vocab_size == 0u) {
        FAIL(AXON_CSYS_BPE_BAD_LAYOUT);
    }
    if ((size_t) pat_len + 20u > blob_len) {
        FAIL(AXON_CSYS_BPE_BAD_LAYOUT);
    }
    size_t pat_off = 16u;
    size_t entries_size_off = pat_off + (size_t) pat_len;
    if (entries_size_off + 4u > blob_len) {
        FAIL(AXON_CSYS_BPE_BAD_LAYOUT);
    }
    uint32_t entries_byte_count = axon_csys_bpe_read_u32_le(blob + entries_size_off);
    size_t entries_off = entries_size_off + 4u;
    if (entries_off + (size_t) entries_byte_count > blob_len) {
        FAIL(AXON_CSYS_BPE_BAD_LAYOUT);
    }

    AxonCsysBpeEncoder* enc = (AxonCsysBpeEncoder*) calloc(1, sizeof *enc);
    if (enc == NULL) {
        FAIL(AXON_CSYS_BPE_OOM);
    }
    enc->blob = blob;
    enc->blob_len = blob_len;
    enc->vocab_size = vocab_size;
    enc->regex_pat = blob + pat_off;
    enc->regex_pat_len = pat_len;

    /* Hash capacity: next_pow2(vocab * 2) with floor at 16. */
    enc->hash_cap = axon_csys_bpe_next_pow2_u32(vocab_size * 2u);
    enc->hash_mask = enc->hash_cap - 1u;
    enc->hash = (uint32_t*) malloc((size_t) enc->hash_cap * sizeof(uint32_t));
    if (enc->hash == NULL) {
        free(enc);
        FAIL(AXON_CSYS_BPE_OOM);
    }
    /* Initialise to empty (UINT32_MAX = 0xFFFFFFFF, all 1-bits). */
    memset(enc->hash, 0xFF, (size_t) enc->hash_cap * sizeof(uint32_t));

    enc->rank_to_offset = (uint32_t*) malloc((size_t) vocab_size * sizeof(uint32_t));
    if (enc->rank_to_offset == NULL) {
        free(enc->hash);
        free(enc);
        FAIL(AXON_CSYS_BPE_OOM);
    }

    /* Walk entries. Each entry: [u8 len][bytes][u32 LE rank]. */
    size_t cursor = entries_off;
    size_t end = entries_off + (size_t) entries_byte_count;
    uint32_t entries_seen = 0u;
    while (cursor < end) {
        if (cursor >= blob_len) {
            free(enc->rank_to_offset);
            free(enc->hash);
            free(enc);
            FAIL(AXON_CSYS_BPE_BAD_LAYOUT);
        }
        uint32_t entry_off = (uint32_t) cursor;
        uint8_t key_len = blob[cursor];
        if (cursor + 1u + (size_t) key_len + 4u > end) {
            free(enc->rank_to_offset);
            free(enc->hash);
            free(enc);
            FAIL(AXON_CSYS_BPE_BAD_LAYOUT);
        }
        uint32_t rank = axon_csys_bpe_read_u32_le(blob + cursor + 1u + (size_t) key_len);
        if (rank >= vocab_size) {
            free(enc->rank_to_offset);
            free(enc->hash);
            free(enc);
            FAIL(AXON_CSYS_BPE_BAD_LAYOUT);
        }
        axon_csys_bpe_hash_insert_unchecked(enc, entry_off);
        enc->rank_to_offset[rank] = entry_off;
        cursor += 1u + (size_t) key_len + 4u;
        ++entries_seen;
    }
    if (entries_seen != vocab_size) {
        free(enc->rank_to_offset);
        free(enc->hash);
        free(enc);
        FAIL(AXON_CSYS_BPE_BAD_LAYOUT);
    }

    if (out_err) {
        *out_err = AXON_CSYS_BPE_OK;
    }
    return enc;

#undef FAIL
}

void axon_csys_bpe_destroy(AxonCsysBpeEncoder* enc) {
    if (enc == NULL) {
        return;
    }
    free(enc->hash);
    free(enc->rank_to_offset);
    free(enc);
}

uint32_t axon_csys_bpe_vocab_size(const AxonCsysBpeEncoder* enc) {
    if (enc == NULL) {
        return 0u;
    }
    return enc->vocab_size;
}

const uint8_t* axon_csys_bpe_regex_pat(
    const AxonCsysBpeEncoder* enc,
    size_t* out_len)
{
    if (enc == NULL) {
        if (out_len) {
            *out_len = 0u;
        }
        return NULL;
    }
    if (out_len) {
        *out_len = (size_t) enc->regex_pat_len;
    }
    return enc->regex_pat;
}

AxonCsysBpeError axon_csys_bpe_token_bytes(
    const AxonCsysBpeEncoder* enc,
    uint32_t rank,
    const uint8_t** out_ptr,
    size_t* out_len)
{
    if (enc == NULL || out_ptr == NULL || out_len == NULL) {
        return AXON_CSYS_BPE_NULL_ARG;
    }
    if (rank >= enc->vocab_size) {
        return AXON_CSYS_BPE_UNKNOWN_RANK;
    }
    uint32_t entry_off = enc->rank_to_offset[rank];
    uint8_t len = enc->blob[entry_off];
    *out_ptr = enc->blob + entry_off + 1u;
    *out_len = (size_t) len;
    return AXON_CSYS_BPE_OK;
}

AxonCsysBpeError axon_csys_bpe_lookup_rank(
    const AxonCsysBpeEncoder* enc,
    const uint8_t* bytes,
    size_t len,
    uint32_t* out_rank)
{
    if (enc == NULL || bytes == NULL || out_rank == NULL) {
        return AXON_CSYS_BPE_NULL_ARG;
    }
    if (len == 0u || len > 255u) {
        return AXON_CSYS_BPE_NOT_FOUND;
    }
    uint32_t bucket = axon_csys_bpe_hash_probe(enc, bytes, len);
    if (bucket == AXON_CSYS_BPE_HASH_EMPTY) {
        return AXON_CSYS_BPE_NOT_FOUND;
    }
    uint32_t entry_off = enc->hash[bucket];
    uint8_t key_len = enc->blob[entry_off];
    *out_rank = axon_csys_bpe_read_u32_le(enc->blob + entry_off + 1u + (size_t) key_len);
    return AXON_CSYS_BPE_OK;
}

/* ──────────────────────────────────────────────────────────────────────
 * 4. byte_pair_merge — canonical tiktoken algorithm.
 *
 * The algorithm operates over an array of `parts`, where each part is
 * a tuple (start_offset, rank_of_pair_starting_here). The pair's rank
 * is the rank of the byte sequence that would result from merging this
 * part with the next part. UINT32_MAX means "no merge possible".
 *
 * Invariant: `parts[len-1]` is a sentinel with start = piece_len. Its
 * rank is irrelevant (no pair starts there). The number of "real"
 * parts is `parts_len - 1`.
 *
 * Iteration:
 *   while min_rank < UINT32_MAX:
 *     i = leftmost index with rank == min_rank
 *     parts[i+1] is removed (memmove down)
 *     parts[i].rank := rank of new merged pair (parts[i].start ..
 *                       parts[i+2].start)  (or UINT32_MAX if no merge)
 *     parts[i-1].rank := same recompute, if i>0
 *     find new min_rank
 *
 * Termination: every iteration removes one part, so this halts in
 * at most `piece_len` iterations.
 *
 * The output is the rank of each remaining real part — looked up by
 * the (start, next_start) byte sequence.
 * ────────────────────────────────────────────────────────────────── */

typedef struct {
    uint32_t start;
    uint32_t rank;
} AxonCsysBpePart;

/* Return the rank of the pair whose merged span is
 * piece[parts[i].start .. parts[i+skip].start], or UINT32_MAX if no
 * such merge exists in the vocabulary. `skip` is 2 in normal usage —
 * we look at parts[i].start..parts[i+2].start, i.e. the merge of
 * parts[i] and parts[i+1]. */
static uint32_t axon_csys_bpe_pair_rank(
    const AxonCsysBpeEncoder* enc,
    const uint8_t* piece,
    const AxonCsysBpePart* parts,
    size_t parts_len,
    size_t i,
    size_t skip)
{
    if (i + skip >= parts_len) {
        return UINT32_MAX;
    }
    uint32_t start = parts[i].start;
    uint32_t end = parts[i + skip].start;
    if (end <= start) {
        return UINT32_MAX;
    }
    size_t span = (size_t) (end - start);
    if (span == 0u || span > 255u) {
        return UINT32_MAX;
    }
    uint32_t bucket = axon_csys_bpe_hash_probe(enc, piece + start, span);
    if (bucket == AXON_CSYS_BPE_HASH_EMPTY) {
        return UINT32_MAX;
    }
    uint32_t entry_off = enc->hash[bucket];
    uint8_t key_len = enc->blob[entry_off];
    return axon_csys_bpe_read_u32_le(enc->blob + entry_off + 1u + (size_t) key_len);
}

AxonCsysBpeError axon_csys_bpe_encode_piece(
    const AxonCsysBpeEncoder* enc,
    const uint8_t* piece,
    size_t piece_len,
    uint32_t* out_ranks,
    size_t out_ranks_capacity,
    size_t* out_count)
{
    if (enc == NULL || (piece == NULL && piece_len > 0u) || out_count == NULL) {
        return AXON_CSYS_BPE_NULL_ARG;
    }
    if (out_ranks == NULL && out_ranks_capacity > 0u) {
        return AXON_CSYS_BPE_NULL_ARG;
    }
    if (piece_len > AXON_CSYS_BPE_MAX_PIECE) {
        return AXON_CSYS_BPE_PIECE_TOO_LONG;
    }

    /* Degenerate cases. */
    if (piece_len == 0u) {
        *out_count = 0u;
        return AXON_CSYS_BPE_OK;
    }
    if (piece_len == 1u) {
        if (out_ranks_capacity < 1u) {
            return AXON_CSYS_BPE_BUFFER_TOO_SMALL;
        }
        uint32_t rank;
        AxonCsysBpeError err = axon_csys_bpe_lookup_rank(enc, piece, 1u, &rank);
        if (err != AXON_CSYS_BPE_OK) {
            return err;
        }
        out_ranks[0] = rank;
        *out_count = 1u;
        return AXON_CSYS_BPE_OK;
    }

    /* parts has piece_len + 1 entries: piece_len real + 1 sentinel. */
    AxonCsysBpePart parts[AXON_CSYS_BPE_MAX_PIECE + 1u];
    size_t parts_len = piece_len + 1u;

    /* Initialise: every byte its own part; ranks are pair-merge ranks. */
    for (size_t i = 0; i < parts_len; ++i) {
        parts[i].start = (uint32_t) i;
        parts[i].rank = UINT32_MAX;
    }
    for (size_t i = 0; i + 2u <= parts_len; ++i) {
        parts[i].rank = axon_csys_bpe_pair_rank(enc, piece, parts, parts_len, i, 2u);
    }
    /* Last real part (parts[parts_len-2]) has no successor pair to merge with;
     * its rank stays UINT32_MAX. The sentinel parts[parts_len-1] also stays. */

    /* Main loop. */
    while (parts_len > 1u) {
        uint32_t min_rank = UINT32_MAX;
        size_t min_idx = 0u;
        for (size_t i = 0; i + 1u < parts_len; ++i) {
            if (parts[i].rank < min_rank) {
                min_rank = parts[i].rank;
                min_idx = i;
            }
        }
        if (min_rank == UINT32_MAX) {
            break;
        }
        /* Merge: remove parts[min_idx + 1]. */
        size_t i = min_idx;
        if (i + 2u < parts_len) {
            memmove(
                &parts[i + 1u],
                &parts[i + 2u],
                (parts_len - i - 2u) * sizeof(AxonCsysBpePart));
        }
        --parts_len;
        /* Recompute parts[i].rank (new merged pair → next part). */
        parts[i].rank = axon_csys_bpe_pair_rank(enc, piece, parts, parts_len, i, 2u);
        /* Recompute parts[i-1].rank (left neighbour's pair rank
         * changed because its right partner widened). */
        if (i > 0u) {
            parts[i - 1u].rank = axon_csys_bpe_pair_rank(enc, piece, parts, parts_len, i - 1u, 2u);
        }
    }

    /* Output: rank of each real part (parts_len - 1 entries). */
    size_t real_count = parts_len - 1u;
    if (out_ranks_capacity < real_count) {
        return AXON_CSYS_BPE_BUFFER_TOO_SMALL;
    }
    for (size_t i = 0; i < real_count; ++i) {
        uint32_t start = parts[i].start;
        uint32_t end = parts[i + 1u].start;
        size_t span = (size_t) (end - start);
        if (span == 0u || span > 255u) {
            return AXON_CSYS_BPE_BAD_LAYOUT;
        }
        uint32_t bucket = axon_csys_bpe_hash_probe(enc, piece + start, span);
        if (bucket == AXON_CSYS_BPE_HASH_EMPTY) {
            /* Surface single-byte gaps with a clearer error.
             * In tiktoken cl100k / o200k every byte 0..255 has its
             * own token, so this can only fire on a corrupt blob. */
            return AXON_CSYS_BPE_NOT_FOUND;
        }
        uint32_t entry_off = enc->hash[bucket];
        uint8_t key_len = enc->blob[entry_off];
        out_ranks[i] = axon_csys_bpe_read_u32_le(enc->blob + entry_off + 1u + (size_t) key_len);
    }
    *out_count = real_count;
    return AXON_CSYS_BPE_OK;
}

/* ──────────────────────────────────────────────────────────────────────
 * 5. UTF-8 byte boundary detection
 *
 * Walks left from `max_offset` until finding a leading byte (or 0).
 * Scalar baseline always available. Lane-acceleration possible on
 * x86_64 (SSE2/AVX2) and aarch64 (NEON) — the baseline is already
 * fast enough for typical (<100 byte) call sites that the SIMD
 * activation is reserved for 25.j benchmarks where input is
 * megabyte-scale.
 *
 * Correctness: a valid UTF-8 character boundary is any byte whose top
 * 2 bits are NOT 10 (i.e. not a continuation byte). This is invariant
 * under all UTF-8 codepoints — ASCII (0xxxxxxx), 2-byte (110xxxxx +
 * 10xxxxxx), 3-byte (1110xxxx + 10..), 4-byte (11110xxx + 10..).
 * Walking left past continuation bytes always lands on a leading
 * byte of a (possibly truncated) codepoint. Per the spec contract
 * (`max_offset == 0 → 0`) we always have a valid stopping point.
 * ────────────────────────────────────────────────────────────────── */

/* Compile-time SIMD capability detection. The `static` impls are
 * always available; the lane-accelerated impl activates only when
 * the compiler can target the relevant ISA. The boundary scan is
 * not hot enough to justify runtime CPU dispatch — adopters who
 * want NEON/AVX builds set `-march=native` or per-package overrides. */
#if defined(__SSE2__) || defined(_M_X64) || defined(_M_AMD64)
#  include <emmintrin.h>
#  define AXON_CSYS_UTF8_HAS_SSE2 1
#endif

#if defined(__ARM_NEON) || defined(__ARM_NEON__)
#  include <arm_neon.h>
#  define AXON_CSYS_UTF8_HAS_NEON 1
#endif

static size_t axon_csys_utf8_boundary_floor_scalar(
    const uint8_t* bytes,
    size_t len,
    size_t max_offset)
{
    /* Position `p` is a valid UTF-8 boundary iff splitting the byte
     * stream at `p` does not break a codepoint. Equivalently:
     *   - p == 0     → boundary (start)
     *   - p == len   → boundary (end)
     *   - 0 < p < len → boundary iff bytes[p] is NOT a continuation
     *                    byte (i.e. its top 2 bits are not `10`).
     *
     * Walk left from `max_offset` while bytes[max_offset] is a
     * continuation byte. The loop runs at most 3 iterations (UTF-8
     * codepoints span at most 4 bytes → at most 3 continuations). */
    if (max_offset >= len) {
        return len;
    }
    while (max_offset > 0u && !axon_csys_utf8_is_leading_byte(bytes[max_offset])) {
        --max_offset;
    }
    return max_offset;
}

size_t axon_csys_utf8_boundary_floor(
    const uint8_t* bytes,
    size_t len,
    size_t max_offset)
{
    if (bytes == NULL || len == 0u || max_offset == 0u) {
        return 0u;
    }
    /* The scalar formulation runs at most 3 iterations (UTF-8 spans
     * at most 4 bytes → at most 3 continuations to walk past). SIMD
     * acceleration on this entry point is not meaningful — adopters
     * who want bulk UTF-8 validation use `axon_csys_utf8_count_chars`
     * instead, which is the genuine SIMD use case.
     *
     * Keeping the function tiny and obviously correct beats a
     * lane-vector dance over a 4-byte window. */
    return axon_csys_utf8_boundary_floor_scalar(bytes, len, max_offset);
}

/* ──────────────────────────────────────────────────────────────────────
 * Bulk UTF-8 character counting — the actual SIMD use case.
 *
 * Counts UTF-8 codepoints in `bytes[0..len]`. A codepoint corresponds
 * to exactly one leading byte (top 2 bits ≠ `10`); count those and
 * we have the codepoint count. SIMD-friendly: 16 bytes at a time on
 * SSE2 / 16 bytes on NEON.
 *
 * Used by token-cost estimation (legacy 4-chars-per-token fallback)
 * and by adopters performing chunk-size budgets in UTF-8 chars rather
 * than bytes.
 * ────────────────────────────────────────────────────────────────── */

#ifdef AXON_CSYS_UTF8_HAS_SSE2

/* Portable popcount over a 16-bit mask. MSVC has `__popcnt` but only
 * with the SSE4.2 extension; gcc / clang have `__builtin_popcount`.
 * SWAR popcount on the 16-bit value: reduce bit pairs → 4-bit groups
 * → 8-bit groups, then sum the two bytes. The final `(m + (m >> 8)) &
 * 0xFFu` step is correct for u16 specifically (the `(m * 0x0101) >> 8`
 * idiom needs an explicit `& 0xFFu` mask to extract just the low byte
 * — a subtle u32-vs-u16 mistake worth flagging here so the bug
 * doesn't return). */
static inline unsigned axon_csys_popcount16(unsigned m) {
    m = m - ((m >> 1) & 0x5555u);
    m = (m & 0x3333u) + ((m >> 2) & 0x3333u);
    m = (m + (m >> 4)) & 0x0F0Fu;
    return (m + (m >> 8)) & 0xFFu;
}

static size_t axon_csys_utf8_count_chars_sse2(
    const uint8_t* bytes,
    size_t len)
{
    /* Process 16 bytes per iteration. Leaders = total - continuations.
     * A continuation byte has top 2 bits == `10` (i.e. b & 0xC0 == 0x80).
     * The vector intrinsic family `_mm_set1_epi8` takes a signed char
     * argument; cast via `int8_t` to get a clean two's-complement
     * reinterpretation of 0xC0 / 0x80 without MSVC's signed-narrowing
     * warning. */
    /* `_mm_set1_epi8` takes a signed-char arg. 0xC0 / 0x80 as two's-
     * complement int8 are -64 / -128. Spelling them this way keeps
     * MSVC /W4 silent (no constant-truncation warning) while leaving
     * the underlying byte pattern unchanged. */
    const __m128i mask_top2 = _mm_set1_epi8(-64);   /* 0xC0 */
    const __m128i mask_cont = _mm_set1_epi8(-128);  /* 0x80 */
    size_t i = 0;
    size_t count = 0;
    for (; i + 16u <= len; i += 16u) {
        /* Load 16 bytes from the byte stream into the SSE register.
         * `_mm_loadu_si128` is the documented unaligned-load
         * intrinsic, but the cast `(const __m128i*)bytes` trips
         * clang's `-Wcast-align` (which scrutinises the cast itself,
         * not the access semantics). `memcpy` into a stack-resident
         * __m128i sidesteps the warning + compiles to the same
         * unaligned MOVDQU instruction at -O1 or higher. */
        __m128i v;
        memcpy(&v, bytes + i, sizeof v);
        __m128i top2 = _mm_and_si128(v, mask_top2);
        __m128i is_cont = _mm_cmpeq_epi8(top2, mask_cont);
        unsigned mask = (unsigned) _mm_movemask_epi8(is_cont);
        unsigned cont_count = axon_csys_popcount16(mask & 0xFFFFu);
        count += (size_t) (16u - cont_count);
    }
    /* Scalar tail. */
    for (; i < len; ++i) {
        if (axon_csys_utf8_is_leading_byte(bytes[i])) {
            ++count;
        }
    }
    return count;
}
#endif

#ifdef AXON_CSYS_UTF8_HAS_NEON
static size_t axon_csys_utf8_count_chars_neon(
    const uint8_t* bytes,
    size_t len)
{
    /* NEON path: 16-byte vector, mask top 2 bits, compare to 0x80,
     * convert to 0/1 mask, sum. */
    const uint8x16_t mask_top2 = vdupq_n_u8(0xC0u);
    const uint8x16_t mask_cont = vdupq_n_u8(0x80u);
    size_t i = 0;
    size_t count = 0;
    for (; i + 16u <= len; i += 16u) {
        uint8x16_t v = vld1q_u8(bytes + i);
        uint8x16_t top2 = vandq_u8(v, mask_top2);
        uint8x16_t is_cont = vceqq_u8(top2, mask_cont);
        /* `is_cont` lanes are 0xFF for continuation bytes else 0x00.
         * Convert to 0x01/0x00 then sum across the vector. */
        uint8x16_t ones = vshrq_n_u8(is_cont, 7);
        uint8_t cont_count = vaddvq_u8(ones);
        count += (size_t) (16u - (size_t) cont_count);
    }
    for (; i < len; ++i) {
        if (axon_csys_utf8_is_leading_byte(bytes[i])) {
            ++count;
        }
    }
    return count;
}
#endif

/* The scalar fallback is selected only when neither SSE2 nor NEON is
 * available — i.e. nowhere in our supported target matrix as of
 * v1.19.x (x86_64 hits SSE2, aarch64 hits NEON). The function still
 * exists for portability + as the canonical reference the SIMD
 * variants must agree with. Annotate so `-Wunused-function` does not
 * fail on platforms that compile only the SIMD variants. */
/* Pre-C23 GCC short-circuit caveat — see probe.c (nested-#ifdef pattern). */
#ifdef __has_c_attribute
#  if __has_c_attribute(maybe_unused)
#    define AXON_CSYS_BPE_MAYBE_UNUSED [[maybe_unused]]
#  endif
#endif
#ifndef AXON_CSYS_BPE_MAYBE_UNUSED
#  if defined(__GNUC__) || defined(__clang__)
#    define AXON_CSYS_BPE_MAYBE_UNUSED __attribute__((unused))
#  else
#    define AXON_CSYS_BPE_MAYBE_UNUSED
#  endif
#endif

AXON_CSYS_BPE_MAYBE_UNUSED
static size_t axon_csys_utf8_count_chars_scalar(
    const uint8_t* bytes,
    size_t len)
{
    size_t count = 0;
    for (size_t i = 0; i < len; ++i) {
        if (axon_csys_utf8_is_leading_byte(bytes[i])) {
            ++count;
        }
    }
    return count;
}

size_t axon_csys_utf8_count_chars(const uint8_t* bytes, size_t len) {
    if (bytes == NULL || len == 0u) {
        return 0u;
    }
#if defined(AXON_CSYS_UTF8_HAS_SSE2)
    return axon_csys_utf8_count_chars_sse2(bytes, len);
#elif defined(AXON_CSYS_UTF8_HAS_NEON)
    return axon_csys_utf8_count_chars_neon(bytes, len);
#else
    return axon_csys_utf8_count_chars_scalar(bytes, len);
#endif
}
