"""Binding over libllama.so - just the calls the entropy sampler needs.

Backed by a cffi extension (see build_capi.py) compiled directly against your
llama.h + libllama.so, so struct layout comes from real compilation instead
of a hand-copied Python guess: a mismatched llama.cpp build fails loudly at
build time rather than silently corrupting memory at run time. Run
`python build_capi.py` once before importing this module.

Not a general-purpose binding - only exposes what fork_peek_test.py / the
sampler use.
"""
import numpy as np

try:
    from _llama_cffi import ffi, lib
except ImportError as e:
    raise RuntimeError(
        "Compiled bindings not found. Run `python build_capi.py` first "
        "(with LLAMA_CPP_LIB, and LLAMA_CPP_INCLUDE if it can't be guessed, set)."
    ) from e


def tokenize(vocab, text, add_special=True, parse_special=True, max_tokens=8192):
    text_bytes = text.encode("utf-8")
    buf = ffi.new("int32_t[]", max_tokens)
    n = lib.llama_tokenize(
        vocab, text_bytes, len(text_bytes), buf, max_tokens, add_special, parse_special,
    )
    if n < 0:
        raise RuntimeError(f"tokenize buffer too small, need {-n}")
    return [buf[i] for i in range(n)]


def token_to_str(vocab, token):
    buf = ffi.new("char[]", 64)
    n = lib.llama_token_to_piece(vocab, token, buf, 64, 0, True)
    if n < 0:
        buf = ffi.new("char[]", -n)
        n = lib.llama_token_to_piece(vocab, token, buf, -n, 0, True)
    return ffi.buffer(buf, n)[:].decode("utf-8", errors="replace")


def as_float_array(ptr, n):
    """Copy an n-element float* returned by the library into a numpy array."""
    return np.frombuffer(ffi.buffer(ptr, n * ffi.sizeof("float")), dtype=np.float32).copy()