"""Byte-level tokenizer shared by questions, choices and text states.

No vocabulary file, no training, works for any language and for arbitrary
choice labels ("UP", "ticket:billing", "negative", ...). ids: 0=pad, 1..3 reserved,
byte b -> b + 4. Pure numpy so the ONNX runtime path needs no torch.
"""
from __future__ import annotations

from typing import Sequence

import numpy as np

PAD = 0
OFFSET = 4
VOCAB_SIZE = 256 + OFFSET


def encode(text: str, max_len: int) -> np.ndarray:
    b = text.encode("utf-8")[:max_len]
    out = np.zeros(max_len, dtype=np.int64)
    if b:
        out[: len(b)] = np.frombuffer(b, dtype=np.uint8).astype(np.int64) + OFFSET
    return out


def encode_batch(texts: Sequence[str], max_len: int) -> np.ndarray:
    out = np.zeros((len(texts), max_len), dtype=np.int64)
    for i, t in enumerate(texts):
        b = t.encode("utf-8")[:max_len]
        if b:
            out[i, : len(b)] = np.frombuffer(b, dtype=np.uint8).astype(np.int64) + OFFSET
    return out


def encode_choice_sets(choice_sets: Sequence[Sequence[str]], max_len: int) -> np.ndarray:
    """(B, Nmax, L) int64; padded choices are all-zero rows (masked by the model)."""
    n = max(len(c) for c in choice_sets)
    out = np.zeros((len(choice_sets), n, max_len), dtype=np.int64)
    for i, cs in enumerate(choice_sets):
        out[i, : len(cs)] = encode_batch(cs, max_len)
    return out