#!/usr/bin/env python3
# Copyright (c) 2023 SpacemiT
import os
import ctypes
import glob

kSpaceMITExecutionProvider = "SpaceMITExecutionProvider"
kExecutionProviderSharedLibraryEntry = "GetSpaceMITSharedProviderFactory"
path_list = glob.glob(
    os.path.join(os.path.dirname(os.path.realpath(__file__)), "libspacemit_ep.so.*")
)
if len(path_list) > 0:
    EPLibPath = path_list[0]
else:
    raise ImportError("can not find libspacemit_ep.so.* in package.")

ep_lib_handle = ctypes.CDLL(EPLibPath)


def __ort_initialize_session_wrap_func(ort_method, *args, **kwargs):
    def _wrap_func(*args, **kwargs):

        sess, providers, provider_options, disabled_optimizers = args[:4]
        assert len(providers) == len(
            provider_options
        ), "providers and provider_options should have same length. "
        if kSpaceMITExecutionProvider in providers:
            ep_index = providers.index(kSpaceMITExecutionProvider)
            provider_options[ep_index]["shared_lib_path"] = EPLibPath
            provider_options[ep_index][
                "provider_factory_entry_point"
            ] = kExecutionProviderSharedLibraryEntry
        if "CPUExecutionProvider" not in providers:
            providers.append("CPUExecutionProvider")
            provider_options.append({})

        return ort_method(*args, **kwargs)

    return _wrap_func


def __ort_get_available_providers_wrap_func(ort_method):
    def _wrap_func():
        eps = ort_method()
        if kSpaceMITExecutionProvider not in eps:
            eps = [kSpaceMITExecutionProvider] + eps
        return eps

    return _wrap_func


def __ort_inferencesession_init_wrap_func(ort_method, *args, **kwargs):
    def _wrap_func(*args, **kwargs):
        session_options = args[1]
        if "get_raw_ptr" not in dir(session_options):
            raise RuntimeError("onnxruntime library maybe mismatch.")
        ep_lib_handle.CreateSpaceMITSessionWrapper(
            ctypes.c_void_p(session_options.get_raw_ptr())
        )
        return ort_method(*args, **kwargs)

    return _wrap_func


def __spacemit_env_init():
    import onnxruntime

    onnxruntime.capi._pybind_state.InferenceSession.initialize_session = (
        __ort_initialize_session_wrap_func(
            onnxruntime.capi._pybind_state.InferenceSession.initialize_session
        )
    )

    onnxruntime.capi._pybind_state.InferenceSession.__init__ = (
        __ort_inferencesession_init_wrap_func(
            onnxruntime.capi._pybind_state.InferenceSession.__init__
        )
    )

    onnxruntime.capi._pybind_state.get_available_providers = (
        __ort_get_available_providers_wrap_func(
            onnxruntime.capi._pybind_state.get_available_providers
        )
    )

    onnxruntime.get_available_providers = __ort_get_available_providers_wrap_func(
        onnxruntime.get_available_providers
    )

    ep_lib_handle.SpaceMITSharedProviderInit()


__spacemit_env_init()
