_api.py

# Copyright 2024-2026 Qualcomm Technologies, Inc. and/or its subsidiaries.
# SPDX-License-Identifier: BSD-3-Clause

"""Bound C functions from libgeniex with argtypes/restype annotations."""

import atexit
import logging
import os
from ctypes import CFUNCTYPE, POINTER, byref, c_char_p, c_int32, c_void_p

from ._lib import load_library
from ._types import (
    geniex_ChipsetList,
    geniex_GetDeviceListInput,
    geniex_GetDeviceListOutput,
    geniex_GetPluginListOutput,
    geniex_HubModelList,
    geniex_KvCacheLoadInput,
    geniex_KvCacheLoadOutput,
    geniex_KvCacheSaveInput,
    geniex_KvCacheSaveOutput,
    geniex_LlmApplyChatTemplateInput,
    geniex_LlmApplyChatTemplateOutput,
    geniex_LlmCreateInput,
    geniex_LlmForwardLogitsInput,
    geniex_LlmForwardLogitsOutput,
    geniex_LlmGenerateInput,
    geniex_LlmGenerateOutput,
    geniex_LlmModelInfo,
    geniex_ModelListDetailedOutput,
    geniex_ModelPaths,
    geniex_ModelPullInput,
    geniex_ModelQueryOutput,
    geniex_ResolveDeviceInput,
    geniex_ResolveDeviceOutput,
    geniex_VlmApplyChatTemplateInput,
    geniex_VlmApplyChatTemplateOutput,
    geniex_VlmCapabilities,
    geniex_VlmCreateInput,
    geniex_VlmGenerateInput,
    geniex_VlmGenerateOutput,
)

# Mirrors geniex_LogLevel in sdk/include/geniex.h — values MUST match.
_LEVEL_TRACE = 0
_LEVEL_DEBUG = 1
_LEVEL_INFO = 2
_LEVEL_WARN = 3
_LEVEL_ERROR = 4
_LEVEL_NONE = 5  # sentinel above ERROR silences SDK output entirely

_LEVEL_STR_TO_PY = {
    'trace': logging.DEBUG,
    'debug': logging.DEBUG,
    'info': logging.INFO,
    'warn': logging.WARNING,
    'error': logging.ERROR,
    'none': logging.CRITICAL + 1,
}

geniex_log_callback = CFUNCTYPE(None, c_int32, c_char_p)

_logger = logging.getLogger('geniex')
_logger.addHandler(logging.NullHandler())

# Strong reference — ctypes callbacks must outlive the C side.
_log_cb_ref: 'geniex_log_callback | None' = None


def _sdk_log_bridge(level: int, msg_bytes):
    try:
        msg = msg_bytes.decode('utf-8', errors='replace') if msg_bytes else ''
    except Exception:
        msg = '<undecodable log message>'
    if level == _LEVEL_ERROR:
        _logger.error(msg)
    elif level == _LEVEL_WARN:
        _logger.warning(msg)
    elif level == _LEVEL_INFO:
        _logger.info(msg)
    else:
        _logger.debug(msg)


def set_log_level(level: str) -> None:
    """Set the ``geniex`` Python logger level.

    Accepts ``'trace' | 'debug' | 'info' | 'warn' | 'error' | 'none'``.
    Unknown values are ignored. Safe to call before or after :func:`init`.
    """
    level_norm = level.lower().strip() if level else ''
    if level_norm not in _LEVEL_STR_TO_PY:
        return
    _logger.setLevel(_LEVEL_STR_TO_PY[level_norm])


def install_log_callback() -> None:
    """Register the SDK → logging bridge and seed the logger from ``GENIEX_LOG``."""
    global _log_cb_ref
    if _log_cb_ref is not None:
        return

    lib = load_library()

    _log_cb_ref = geniex_log_callback(_sdk_log_bridge)
    lib.geniex_set_log(_log_cb_ref)

    requested = os.environ.get('GENIEX_LOG', '').strip().lower()
    if requested in {'off', '0', 'false'}:
        requested = 'none'
    if requested in _LEVEL_STR_TO_PY:
        _logger.setLevel(_LEVEL_STR_TO_PY[requested])


class GenieXError(Exception):
    """Raised when a geniex C call returns a negative error code."""

    def __init__(self, code: int, message: str):
        super().__init__(f'GenieXError({code}): {message}')
        self.code = code


# Subset of geniex_ErrorCode values that callers may want to check against.
# Keep aligned with sdk/include/geniex.h; expand on demand.
GENIEX_ERROR_COMMON_PLUGIN_INVALID = -100302
GENIEX_ERROR_LLM_TOKENIZATION_CONTEXT_LENGTH = -200004


def _unknown_runtime_message(runtime: str, available: list[str]) -> str:
    return f"Unknown runtime: {runtime}. Available runtimes: {', '.join(sorted(available))}."


def _check(code: int) -> None:
    if code < 0:
        lib = load_library()
        msg_bytes = lib.geniex_get_error_message(c_int32(code))
        msg = msg_bytes.decode() if msg_bytes else 'unknown error'
        raise GenieXError(code, msg)


def _bind_all() -> None:
    lib = load_library()

    lib.geniex_get_error_message.argtypes = [c_int32]
    lib.geniex_get_error_message.restype = c_char_p

    lib.geniex_set_log.argtypes = [geniex_log_callback]
    lib.geniex_set_log.restype = c_int32

    lib.geniex_init.argtypes = []
    lib.geniex_init.restype = c_int32

    lib.geniex_deinit.argtypes = []
    lib.geniex_deinit.restype = c_int32

    lib.geniex_free.argtypes = [c_void_p]
    lib.geniex_free.restype = None

    lib.geniex_version.argtypes = []
    lib.geniex_version.restype = c_char_p

    lib.geniex_get_plugin_version.argtypes = [c_char_p]
    lib.geniex_get_plugin_version.restype = c_char_p

    lib.geniex_get_plugin_list.argtypes = [POINTER(geniex_GetPluginListOutput)]
    lib.geniex_get_plugin_list.restype = c_int32

    lib.geniex_get_device_list.argtypes = [POINTER(geniex_GetDeviceListInput), POINTER(geniex_GetDeviceListOutput)]
    lib.geniex_get_device_list.restype = c_int32

    lib.geniex_resolve_device.argtypes = [POINTER(geniex_ResolveDeviceInput), POINTER(geniex_ResolveDeviceOutput)]
    lib.geniex_resolve_device.restype = c_int32

    # LLM
    lib.geniex_llm_create.argtypes = [POINTER(geniex_LlmCreateInput), POINTER(c_void_p)]
    lib.geniex_llm_create.restype = c_int32

    lib.geniex_llm_destroy.argtypes = [c_void_p]
    lib.geniex_llm_destroy.restype = c_int32

    lib.geniex_llm_reset.argtypes = [c_void_p]
    lib.geniex_llm_reset.restype = c_int32

    lib.geniex_llm_generate.argtypes = [c_void_p, POINTER(geniex_LlmGenerateInput), POINTER(geniex_LlmGenerateOutput)]
    lib.geniex_llm_generate.restype = c_int32

    lib.geniex_llm_get_model_info.argtypes = [c_void_p, POINTER(geniex_LlmModelInfo)]
    lib.geniex_llm_get_model_info.restype = c_int32

    lib.geniex_llm_forward_logits.argtypes = [
        c_void_p,
        POINTER(geniex_LlmForwardLogitsInput),
        POINTER(geniex_LlmForwardLogitsOutput),
    ]
    lib.geniex_llm_forward_logits.restype = c_int32

    lib.geniex_llm_apply_chat_template.argtypes = [
        c_void_p,
        POINTER(geniex_LlmApplyChatTemplateInput),
        POINTER(geniex_LlmApplyChatTemplateOutput),
    ]
    lib.geniex_llm_apply_chat_template.restype = c_int32

    lib.geniex_llm_save_kv_cache.argtypes = [
        c_void_p,
        POINTER(geniex_KvCacheSaveInput),
        POINTER(geniex_KvCacheSaveOutput),
    ]
    lib.geniex_llm_save_kv_cache.restype = c_int32

    lib.geniex_llm_load_kv_cache.argtypes = [
        c_void_p,
        POINTER(geniex_KvCacheLoadInput),
        POINTER(geniex_KvCacheLoadOutput),
    ]
    lib.geniex_llm_load_kv_cache.restype = c_int32

    # VLM
    lib.geniex_vlm_create.argtypes = [POINTER(geniex_VlmCreateInput), POINTER(c_void_p)]
    lib.geniex_vlm_create.restype = c_int32

    lib.geniex_vlm_destroy.argtypes = [c_void_p]
    lib.geniex_vlm_destroy.restype = c_int32

    lib.geniex_vlm_reset.argtypes = [c_void_p]
    lib.geniex_vlm_reset.restype = c_int32

    lib.geniex_vlm_generate.argtypes = [c_void_p, POINTER(geniex_VlmGenerateInput), POINTER(geniex_VlmGenerateOutput)]
    lib.geniex_vlm_generate.restype = c_int32

    lib.geniex_vlm_apply_chat_template.argtypes = [
        c_void_p,
        POINTER(geniex_VlmApplyChatTemplateInput),
        POINTER(geniex_VlmApplyChatTemplateOutput),
    ]
    lib.geniex_vlm_apply_chat_template.restype = c_int32

    lib.geniex_vlm_get_capabilities.argtypes = [c_void_p, POINTER(geniex_VlmCapabilities)]
    lib.geniex_vlm_get_capabilities.restype = c_int32

    # Model manager
    lib.geniex_model_init.argtypes = [c_char_p]
    lib.geniex_model_init.restype = c_int32

    lib.geniex_model_deinit.argtypes = []
    lib.geniex_model_deinit.restype = c_int32

    lib.geniex_model_pull.argtypes = [POINTER(geniex_ModelPullInput)]
    lib.geniex_model_pull.restype = c_int32

    lib.geniex_model_last_error_message.argtypes = []
    lib.geniex_model_last_error_message.restype = c_char_p

    lib.geniex_model_remove.argtypes = [c_char_p]
    lib.geniex_model_remove.restype = c_int32

    lib.geniex_model_clean.argtypes = [POINTER(c_int32)]
    lib.geniex_model_clean.restype = c_int32

    lib.geniex_model_get_paths.argtypes = [c_char_p, POINTER(geniex_ModelPaths)]
    lib.geniex_model_get_paths.restype = c_int32

    lib.geniex_model_paths_free.argtypes = [POINTER(geniex_ModelPaths)]
    lib.geniex_model_paths_free.restype = None

    lib.geniex_model_get_type.argtypes = [c_char_p, POINTER(c_int32)]
    lib.geniex_model_get_type.restype = c_int32

    lib.geniex_model_set_type.argtypes = [c_char_p, c_int32]
    lib.geniex_model_set_type.restype = c_int32

    lib.geniex_model_list_detailed.argtypes = [POINTER(geniex_ModelListDetailedOutput)]
    lib.geniex_model_list_detailed.restype = c_int32

    lib.geniex_model_list_detailed_free.argtypes = [POINTER(geniex_ModelListDetailedOutput)]
    lib.geniex_model_list_detailed_free.restype = None

    lib.geniex_model_query.argtypes = [
        POINTER(geniex_ModelPullInput),
        POINTER(geniex_ModelQueryOutput),
    ]
    lib.geniex_model_query.restype = c_int32

    lib.geniex_model_query_free.argtypes = [POINTER(geniex_ModelQueryOutput)]
    lib.geniex_model_query_free.restype = None

    lib.geniex_model_resolve_alias.argtypes = [c_char_p, POINTER(c_char_p)]
    lib.geniex_model_resolve_alias.restype = c_int32

    lib.geniex_model_resolve_hub.argtypes = [c_char_p, c_int32, POINTER(c_int32)]
    lib.geniex_model_resolve_hub.restype = c_int32

    lib.geniex_model_list_chipsets.argtypes = [POINTER(geniex_ChipsetList)]
    lib.geniex_model_list_chipsets.restype = c_int32

    lib.geniex_model_list_chipsets_free.argtypes = [POINTER(geniex_ChipsetList)]
    lib.geniex_model_list_chipsets_free.restype = None

    lib.geniex_model_detect_chipset.argtypes = [c_int32, POINTER(c_char_p)]
    lib.geniex_model_detect_chipset.restype = c_int32

    lib.geniex_model_list_hub.argtypes = [c_char_p, POINTER(geniex_HubModelList)]
    lib.geniex_model_list_hub.restype = c_int32

    lib.geniex_model_list_hub_free.argtypes = [POINTER(geniex_HubModelList)]
    lib.geniex_model_list_hub_free.restype = None


_bound = False


def _ensure_bound() -> None:
    global _bound
    if not _bound:
        _bind_all()
        _bound = True
        install_log_callback()


_initialized = False


def init() -> None:
    """Initialise the geniex runtime. Idempotent."""
    global _initialized
    if _initialized:
        return
    _ensure_bound()
    lib = load_library()
    _check(lib.geniex_init())
    _initialized = True
    atexit.register(deinit)


def deinit() -> None:
    """Tear down the geniex runtime."""
    global _initialized
    if not _initialized:
        return
    lib = load_library()
    lib.geniex_deinit()
    _initialized = False


def ensure_init() -> None:
    if not _initialized:
        init()


def _require_init(func_name: str) -> None:
    # Plugin/device queries silently returned empty before init (#715).
    # Surface the missing init step explicitly instead.
    if not _initialized:
        raise RuntimeError(f'geniex.{func_name}() requires the runtime to be initialized. Call geniex.init() first.')


def _encode(s: str | None) -> bytes | None:
    return s.encode() if s else None


def _str_list_to_c(strings: list[str]):
    count = len(strings)
    if count == 0:
        return None, 0
    arr_type = c_char_p * count
    arr = arr_type(*[s.encode() for s in strings])
    return arr, count


def get_runtime_list() -> list[str]:
    """Return the runtime ids registered with libgeniex.

    Raises :class:`RuntimeError` if :func:`init` has not been called — the
    runtime registry is populated by ``geniex_init`` and would otherwise
    silently return ``[]``.
    """
    _require_init('get_runtime_list')
    _ensure_bound()
    lib = load_library()
    out = geniex_GetPluginListOutput()
    _check(lib.geniex_get_plugin_list(byref(out)))
    result = [out.plugin_ids[i].decode() for i in range(out.plugin_count)]
    if out.plugin_ids:
        lib.geniex_free(out.plugin_ids)
    return result


def get_compute_unit_list(runtime: str) -> list[tuple[str, str]]:
    """Return ``[(compute_unit, compute_unit_name), ...]`` for ``runtime``.

    Raises :class:`RuntimeError` if :func:`init` has not been called.
    """
    _require_init('get_compute_unit_list')
    _ensure_bound()
    available = get_runtime_list()
    if runtime not in available:
        raise GenieXError(GENIEX_ERROR_COMMON_PLUGIN_INVALID, _unknown_runtime_message(runtime, available))
    lib = load_library()
    inp = geniex_GetDeviceListInput(plugin_id=runtime.encode())
    out = geniex_GetDeviceListOutput()
    _check(lib.geniex_get_device_list(byref(inp), byref(out)))
    result = [(out.device_ids[i].decode(), out.device_names[i].decode()) for i in range(out.device_count)]
    if out.device_ids:
        lib.geniex_free(out.device_ids)
    if out.device_names:
        lib.geniex_free(out.device_names)
    return result


def resolve_device(
    plugin_id: str,
    model_name: str | None,
    mode: str | None,
    ngl_default: int,
) -> tuple[str | None, int, str | None]:
    """Raw SDK device-alias resolver. Prefer :func:`geniex.resolve_device_map`.

    Returns ``(device_id, ngl, warning)``; ``device_id`` / ``warning`` may
    be ``None``. Raises :class:`GenieXError` for unknown aliases.
    """
    _ensure_bound()
    lib = load_library()
    inp = geniex_ResolveDeviceInput(
        plugin_id=plugin_id.encode(),
        model_name=model_name.encode() if model_name else None,
        mode=mode.encode() if mode else None,
        ngl_default=int(ngl_default),
    )
    out = geniex_ResolveDeviceOutput()
    _check(lib.geniex_resolve_device(byref(inp), byref(out)))

    device_id: str | None = None
    warning: str | None = None
    if out.device_id:
        device_id = c_char_p(out.device_id).value.decode()
        lib.geniex_free(out.device_id)
    if out.warning:
        warning = c_char_p(out.warning).value.decode()
        lib.geniex_free(out.warning)
    return device_id, int(out.ngl), warning


def version() -> str:
    """Return the geniex SDK version string."""
    _ensure_bound()
    lib = load_library()
    return lib.geniex_version().decode()


def get_plugin_version(plugin_id: str) -> str:
    """Return the version string the named plugin reports for itself.

    Plugin must be registered (i.e. :func:`init` has been called and the
    plugin's shared library was discoverable). Returns the empty string
    when the plugin is unknown.
    """
    ensure_init()
    lib = load_library()
    raw = lib.geniex_get_plugin_version(plugin_id.encode())
    return raw.decode() if raw else ''