"""ctypes bridge to the GGML C++ inference engine (liblocalvqe.so).

This is the same shared library production users build from the
inference repo (`-DLOCALVQE_BUILD_SHARED=ON`); the demo loads released
.gguf files through it so listeners hear the deployed artifact, not a
PyTorch re-implementation. The bundle in ./lib is produced by
build_lib.sh (portable build, per-ISA ggml CPU backends dispatched at
runtime).

Each GGUFModel owns one streaming context; a lock serialises requests
(the context is stateful) and `localvqe_reset` runs before every clip
so each request starts from clean stream state.
"""
import ctypes
import os
import threading
from pathlib import Path

import numpy as np

_LIB = None
_LIB_ERR = None
_F32P = ctypes.POINTER(ctypes.c_float)


def _lib_dir() -> Path:
    v = os.environ.get("LOCALVQE_LIB_DIR")
    return Path(v) if v else Path(__file__).resolve().parent / "lib"


def _load_lib():
    global _LIB, _LIB_ERR
    if _LIB is not None or _LIB_ERR is not None:
        return _LIB
    d = _lib_dir()
    try:
        # Dependencies first, RTLD_GLOBAL, so liblocalvqe resolves them
        # without rpath help; ggml then self-resolves its per-ISA CPU
        # backend .so files from the same directory.
        for dep in ("libggml-base.so", "libggml.so"):
            p = d / dep
            if p.exists():
                ctypes.CDLL(str(p), mode=ctypes.RTLD_GLOBAL)
        lib = ctypes.CDLL(str(d / "liblocalvqe.so"), mode=ctypes.RTLD_GLOBAL)
        lib.localvqe_new.argtypes = [ctypes.c_char_p]
        lib.localvqe_new.restype = ctypes.c_size_t
        lib.localvqe_free.argtypes = [ctypes.c_size_t]
        lib.localvqe_reset.argtypes = [ctypes.c_size_t]
        lib.localvqe_process_f32.argtypes = [
            ctypes.c_size_t, _F32P, _F32P, ctypes.c_int, _F32P]
        lib.localvqe_process_f32.restype = ctypes.c_int
        lib.localvqe_last_error.argtypes = [ctypes.c_size_t]
        lib.localvqe_last_error.restype = ctypes.c_char_p
        _LIB = lib
    except OSError as e:
        _LIB_ERR = str(e)
        print(f"GGML engine unavailable ({d}): {e}")
    return _LIB


class GGUFModel:
    """One loaded .gguf behind the C API. process() is clip-at-a-time."""

    def __init__(self, gguf_path: str):
        lib = _load_lib()
        if lib is None:
            raise RuntimeError(f"liblocalvqe.so not loadable: {_LIB_ERR}")
        self._lib = lib
        self._ctx = lib.localvqe_new(str(gguf_path).encode())
        if not self._ctx:
            raise RuntimeError(f"localvqe_new failed for {gguf_path}")
        self.path = str(gguf_path)
        self._lock = threading.Lock()

    def process(self, mic: np.ndarray, ref: np.ndarray) -> np.ndarray:
        n = max(len(mic), len(ref), 512)
        m = np.zeros(n, dtype=np.float32)
        r = np.zeros(n, dtype=np.float32)
        m[:len(mic)] = mic
        r[:len(ref)] = ref
        out = np.zeros(n, dtype=np.float32)
        with self._lock:
            self._lib.localvqe_reset(self._ctx)
            ret = self._lib.localvqe_process_f32(
                self._ctx,
                m.ctypes.data_as(_F32P), r.ctypes.data_as(_F32P),
                ctypes.c_int(n), out.ctypes.data_as(_F32P))
        if ret != 0:
            err = self._lib.localvqe_last_error(self._ctx)
            raise RuntimeError(f"localvqe_process_f32: "
                               f"{err.decode() if err else ret}")
        return out[:max(len(mic), len(ref))]

    def __del__(self):
        ctx = getattr(self, "_ctx", 0)
        if ctx and getattr(self, "_lib", None):
            self._lib.localvqe_free(ctx)
            self._ctx = 0