"""Load MLC LLM library and _ffi_api functions."""

import ctypes
import os
import sys

import tvm
import tvm.base

from . import libinfo

SKIP_LOADING_MLCLLM_SO = os.environ.get("SKIP_LOADING_MLCLLM_SO", "0")


def _load_mlc_llm_lib():
    """Load MLC LLM lib"""
    if sys.platform.startswith("win32") and sys.version_info >= (3, 8):
        for path in libinfo.get_dll_directories():
            os.add_dll_directory(path)
    lib_name = "mlc_llm" if tvm.base._RUNTIME_ONLY else "mlc_llm_module"
    lib_path = libinfo.find_lib_path(lib_name, optional=False)
    return ctypes.CDLL(lib_path[0]), lib_path[0]


@tvm.register_global_func("mlc.debug_cuda_profiler_start")
def _debug_cuda_profiler_start() -> None:
    """Start cuda profiler."""
    import cuda
    import cuda.cudart

    cuda.cudart.cudaProfilerStart()


@tvm.register_global_func("mlc.debug_cuda_profiler_stop")
def _debug_cuda_profiler_stop() -> None:
    """Stop cuda profiler."""
    import cuda
    import cuda.cudart

    cuda.cudart.cudaProfilerStop()


# only load once here
if SKIP_LOADING_MLCLLM_SO == "0":
    _LIB, _LIB_PATH = _load_mlc_llm_lib()