#ifndef PROTON_DRIVER_DISPATCH_H_
#define PROTON_DRIVER_DISPATCH_H_

#include <dlfcn.h>

#include "Utility/Env.h"
#include <stdexcept>
#include <string>

#define DISPATCH_ARGS_0()
#define DISPATCH_ARGS_1(t1) t1 v1
#define DISPATCH_ARGS_2(t1, t2) t1 v1, t2 v2
#define DISPATCH_ARGS_3(t1, t2, t3) t1 v1, t2 v2, t3 v3
#define DISPATCH_ARGS_4(t1, t2, t3, t4) t1 v1, t2 v2, t3 v3, t4 v4
#define DISPATCH_ARGS_5(t1, t2, t3, t4, t5) t1 v1, t2 v2, t3 v3, t4 v4, t5 v5
#define DISPATCH_ARGS_6(t1, t2, t3, t4, t5, t6)                                \
  t1 v1, t2 v2, t3 v3, t4 v4, t5 v5, t6 v6
#define DISPATCH_ARGS_7(t1, t2, t3, t4, t5, t6, t7)                            \
  t1 v1, t2 v2, t3 v3, t4 v4, t5 v5, t6 v6, t7 v7
#define DISPATCH_ARGS_8(t1, t2, t3, t4, t5, t6, t7, t8)                        \
  t1 v1, t2 v2, t3 v3, t4 v4, t5 v5, t6 v6, t7 v7, t8 v8
#define DISPATCH_ARGS_9(t1, t2, t3, t4, t5, t6, t7, t8, t9)                    \
  t1 v1, t2 v2, t3 v3, t4 v4, t5 v5, t6 v6, t7 v7, t8 v8, t9 v9
#define DISPATCH_ARGS_10(t1, t2, t3, t4, t5, t6, t7, t8, t9, t10)              \
  t1 v1, t2 v2, t3 v3, t4 v4, t5 v5, t6 v6, t7 v7, t8 v8, t9 v9, t10 v10
#define DISPATCH_ARGS_11(t1, t2, t3, t4, t5, t6, t7, t8, t9, t10, t11)         \
  t1 v1, t2 v2, t3 v3, t4 v4, t5 v5, t6 v6, t7 v7, t8 v8, t9 v9, t10 v10,      \
      t11 v11
#define DISPATCH_ARGS_12(t1, t2, t3, t4, t5, t6, t7, t8, t9, t10, t11, t12)    \
  t1 v1, t2 v2, t3 v3, t4 v4, t5 v5, t6 v6, t7 v7, t8 v8, t9 v9, t10 v10,      \
      t11 v11, t12 v12
#define DISPATCH_ARGS_N(_12, _11, _10, _9, _8, _7, _6, _5, _4, _3, _2, _1, _0, \
                        N, ...)                                                \
  DISPATCH_ARGS##N
#define DISPATCH_ARGS(...)                                                     \
  DISPATCH_ARGS_N(_0, ##__VA_ARGS__, _12, _11, _10, _9, _8, _7, _6, _5, _4,    \
                  _3, _2, _1, _0)(__VA_ARGS__)

#define DISPATCH_VALS_0()
#define DISPATCH_VALS_1(t1) , v1
#define DISPATCH_VALS_2(t1, t2) , v1, v2
#define DISPATCH_VALS_3(t1, t2, t3) , v1, v2, v3
#define DISPATCH_VALS_4(t1, t2, t3, t4) , v1, v2, v3, v4
#define DISPATCH_VALS_5(t1, t2, t3, t4, t5) , v1, v2, v3, v4, v5
#define DISPATCH_VALS_6(t1, t2, t3, t4, t5, t6) , v1, v2, v3, v4, v5, v6
#define DISPATCH_VALS_7(t1, t2, t3, t4, t5, t6, t7) , v1, v2, v3, v4, v5, v6, v7
#define DISPATCH_VALS_8(t1, t2, t3, t4, t5, t6, t7, t8)                        \
  , v1, v2, v3, v4, v5, v6, v7, v8
#define DISPATCH_VALS_9(t1, t2, t3, t4, t5, t6, t7, t8, t9)                    \
  , v1, v2, v3, v4, v5, v6, v7, v8, v9
#define DISPATCH_VALS_10(t1, t2, t3, t4, t5, t6, t7, t8, t9, t10)              \
  , v1, v2, v3, v4, v5, v6, v7, v8, v9, v10
#define DISPATCH_VALS_11(t1, t2, t3, t4, t5, t6, t7, t8, t9, t10, t11)         \
  , v1, v2, v3, v4, v5, v6, v7, v8, v9, v10, v11
#define DISPATCH_VALS_12(t1, t2, t3, t4, t5, t6, t7, t8, t9, t10, t11, t12)    \
  , v1, v2, v3, v4, v5, v6, v7, v8, v9, v10, v11, v12
#define DISPATCH_VALS_N(_12, _11, _10, _9, _8, _7, _6, _5, _4, _3, _2, _1, _0, \
                        N, ...)                                                \
  DISPATCH_VALS##N
#define DISPATCH_VALS(...)                                                     \
  DISPATCH_VALS_N(_0, ##__VA_ARGS__, _12, _11, _10, _9, _8, _7, _6, _5, _4,    \
                  _3, _2, _1, _0)(__VA_ARGS__)

#define DEFINE_DISPATCH_TEMPLATE(CheckSuccess, FuncName, ExternLib, FuncType,  \
                                 ...)                                          \
  template <>                                                                  \
  ExternLib::RetType FuncName<CheckSuccess>(DISPATCH_ARGS(__VA_ARGS__)) {      \
    typedef typename ExternLib::RetType (*FuncType##_t)(__VA_ARGS__);          \
    static FuncType##_t func = nullptr;                                        \
    return Dispatch<ExternLib>::exec<CheckSuccess, FuncType##_t>(              \
        func, #FuncType DISPATCH_VALS(__VA_ARGS__));                           \
  }

#define DEFINE_DISPATCH(ExternLib, FuncName, FuncType, ...)                    \
  DEFINE_DISPATCH_TEMPLATE(true, FuncName, ExternLib, FuncType, __VA_ARGS__)   \
  DEFINE_DISPATCH_TEMPLATE(false, FuncName, ExternLib, FuncType, __VA_ARGS__)

namespace proton {

struct ExternLibBase {
  using RetType = int; // Generic type, can be overridden in derived structs
  static constexpr const char *name = "";    // Placeholder
  static constexpr const char *symbolName{}; // Placeholder
  static constexpr const char *pathEnv{};    // Placeholder
  static constexpr RetType success = 0;      // Placeholder
  ExternLibBase() = delete;
  ExternLibBase(const ExternLibBase &) = delete;
  ExternLibBase &operator=(const ExternLibBase &) = delete;
  static inline void *lib{nullptr};
};

template <typename ExternLib> class Dispatch {
public:
  Dispatch() = delete;

  static void init(const char *name, void **lib) {
    if (*lib == nullptr) {
      // If not found, try to load it from the default path
      auto dir =
          ExternLib::pathEnv == nullptr ? "" : getStrEnv(ExternLib::pathEnv);
      if (!dir.empty()) {
        auto fullPath = dir + "/" + name;
        *lib = dlopen(fullPath.c_str(), RTLD_LOCAL | RTLD_LAZY);
      } else {
        // Only if the default path is not set, we try to load it from the
        // system.
        // First reuse the existing handle
        *lib = dlopen(name, RTLD_NOLOAD);
        if (*lib == nullptr) {
          // If not found, try to load it from LD_LIBRARY_PATH
          *lib = dlopen(name, RTLD_LOCAL | RTLD_LAZY);
        }
      }
    }
    if (*lib == nullptr) {
      throw std::runtime_error("Could not load `" + std::string(name) + "`");
    }
  }

  static void check(typename ExternLib::RetType ret, const char *functionName) {
    if (ret != ExternLib::success) {
      throw std::runtime_error("Failed to execute " +
                               std::string(functionName) + " with error " +
                               std::to_string(ret));
    }
  }

  template <bool CheckSuccess, typename FnT, typename... Args>
  static typename ExternLib::RetType
  exec(FnT &handler, const char *functionName, Args... args) {
    init(ExternLib::name, &ExternLib::lib);
    if (handler == nullptr) {
      handler = reinterpret_cast<FnT>(dlsym(ExternLib::lib, functionName));
      if (handler == nullptr) {
        throw std::runtime_error("Failed to load " +
                                 std::string(ExternLib::name));
      }
    }
    auto ret = handler(args...);
    if constexpr (CheckSuccess) {
      check(ret, functionName);
    }
    return ret;
  }

  static std::string getLibPath() {
    if (ExternLib::lib == nullptr) {
      // Force initialization
      Dispatch<ExternLib>::init(ExternLib::name, &ExternLib::lib);
      if (ExternLib::lib == nullptr) {
        return "";
      }
    }
    if (ExternLib::lib != nullptr) {
      void *sym = dlsym(ExternLib::lib,
                        ExternLib::symbolName); // pick any known symbol
      Dl_info info;
      if (dladdr(sym, &info)) {
        return info.dli_fname;
      }
    }
    return "";
  }
};

} // namespace proton

#endif // PROTON_DRIVER_DISPATCH_H_
