"""Turn (text, candidate span) into a fixed-size byte window for the model."""
import numpy as np

LEFT, CAND, RIGHT = 64, 96, 32
SEQ = LEFT + CAND + RIGHT + 2
PAD, SEP1, SEP2 = 0, 1, 2  # bytes 0..2 are reserved; real bytes are shifted by +3
VOCAB = 256 + 3
TYPES = ["none", "password", "api_key", "token", "private_key", "generic"]


def _b(s):
    return np.frombuffer(s.encode("utf-8", "replace"), dtype=np.uint8).astype(np.int64) + 3


def encode(text, start, end):
    """Return (ids[SEQ], seg[SEQ]) where seg is 0 pad, 1 left ctx, 2 candidate, 3 right ctx."""
    left = _b(text[max(0, start - LEFT * 2):start])[-LEFT:]
    cand = _b(text[start:end])
    if len(cand) > CAND:  # keep head and tail of very long values (e.g. PEM blocks)
        cand = np.concatenate([cand[:CAND // 2], cand[-CAND // 2:]])
    right = _b(text[end:end + RIGHT * 2])[:RIGHT]
    ids = np.zeros(SEQ, dtype=np.int64)
    seg = np.zeros(SEQ, dtype=np.int64)
    # left context is right-aligned so the candidate always starts at the same offset
    ids[LEFT - len(left):LEFT] = left; seg[LEFT - len(left):LEFT] = 1
    ids[LEFT] = SEP1; seg[LEFT] = 2
    o = LEFT + 1
    ids[o:o + len(cand)] = cand; seg[o:o + len(cand)] = 2
    o += len(cand)
    ids[o] = SEP2; seg[o] = 2
    o += 1
    ids[o:o + len(right)] = right; seg[o:o + len(right)] = 3
    return ids, seg


def encode_string(value):
    """Classify a bare string with no context."""
    return encode(value, 0, len(value))


def label_candidate(cs, ce, spans):
    """Label a candidate against ground-truth spans [(s, e, label, kind)].

    Returns (label, kind) or None when ambiguous (partial overlap) so the sample is dropped.
    """
    clen = ce - cs
    for s, e, lab, kind in spans:
        inter = min(ce, e) - max(cs, s)
        if inter <= 0:
            continue
        if lab == 1:
            if cs <= s and ce >= e and clen <= (e - s) + 40:
                return 1, kind  # candidate contains the secret (e.g. 'Bearer xxx', URL with creds)
            if inter >= 0.8 * (e - s) and inter >= 0.6 * clen:
                return 1, kind
            return None
        else:
            if inter >= 0.6 * clen:
                return 0, kind
            return None
    return 0, "other"