// Copyright (c) Microsoft Corporation. All rights reserved.
// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// Licensed under the MIT License.

#include <unordered_set>

#include "core/framework/provider_options.h"
#include "nv_execution_provider_custom_ops.h"
#include "nv_execution_provider.h"
#include "nv_platform_utils.h"

namespace onnxruntime {
extern TensorrtLogger& GetTensorrtLogger(bool verbose);

/*
 * Create custom op domain list for TRT plugins.
 *
 * Here, we collect all registered TRT plugins from TRT registry and create custom ops with "trt.plugins" domain.
 * Additionally, if users specify extra plugin libraries, TRT EP will load them at runtime which will register those
 * plugins to TRT plugin registry and later TRT EP can get them as well.
 *
 * There are several TRT plugins registered as onnx schema op through contrib op with ONNX domain in the past,
 * for example, EfficientNMS_TRT, MultilevelCropAndResize_TRT, PyramidROIAlign_TRT and DisentangledAttention_TRT.
 * In order not to break the old models using those TRT plugins which were registered with ONNX domain and maintain
 * backward compatible, we need to keep those legacy TRT plugins registered with ONNX domain with contrib ops.
 *
 * Note: Current TRT plugin doesn't have APIs to get number of inputs/outputs of the plugin.
 * So, TensorRTCustomOp uses variadic inputs/outputs to pass ONNX graph validation.
 */
common::Status CreateTensorRTCustomOpDomainList(std::vector<OrtCustomOpDomain*>& domain_list, const std::string extra_plugin_lib_paths) {
  // Domain for TRT plugin custom ops (domain name: "trt.plugins"). Owns the OrtCustomOpDomain object.
  // Raw pointers from .get() are handed out to callers via domain_list and may be held by InferenceSession.
  static std::unique_ptr<OrtCustomOpDomain> custom_op_domain = std::make_unique<OrtCustomOpDomain>();

  // Owns the TensorRTCustomOp objects for TRT plugins. Raw pointers are stored in custom_op_domain->custom_ops_.
  static std::vector<std::unique_ptr<TensorRTCustomOp>> created_custom_op_list;

  // Domain for native custom ops (domain name: "trt"). Owns the OrtCustomOpDomain object.
  // Raw pointers from .get() are handed out to callers via domain_list and may be held by InferenceSession.
  static std::unique_ptr<OrtCustomOpDomain> native_custom_op_domain = std::make_unique<OrtCustomOpDomain>();

  // Owns the TensorRTCustomOp objects for native custom ops. Raw pointers are stored in native_custom_op_domain->custom_ops_.
  // Non-empty list indicates native custom ops have been registered (used to avoid re-registration on subsequent calls).
  static std::vector<std::unique_ptr<TensorRTCustomOp>> native_custom_op_list;

  // Protects concurrent access to all the above static members.
  static std::mutex mutex;
  std::lock_guard<std::mutex> lock(mutex);

  // Add already-initialized native ops to domain list
  if (!native_custom_op_list.empty()) {
    domain_list.push_back(native_custom_op_domain.get());
  }

  if (custom_op_domain->domain_ != "" && custom_op_domain->custom_ops_.size() > 0) {
    domain_list.push_back(custom_op_domain.get());
    return Status::OK();
  }

  // Load any extra TRT plugin library if any.
  // When the TRT plugin library is loaded, the global static object is created and the plugin is registered to TRT registry.
  // This is done through macro, for example, REGISTER_TENSORRT_PLUGIN(VisionTransformerPluginCreator).
  // extra_plugin_lib_paths has the format of "path_1;path_2....;path_n"
  static bool is_loaded = false;
  if (!extra_plugin_lib_paths.empty() && !is_loaded) {
    std::stringstream extra_plugin_libs(extra_plugin_lib_paths);
    std::string lib;
    while (std::getline(extra_plugin_libs, lib, ';')) {
      auto status = LoadDynamicLibrary(ToPathString(lib));
      if (status == Status::OK()) {
        LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Successfully load " << lib;
      } else {
        LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP]" << status.ToString();
      }
    }
    is_loaded = true;
  }

  // Load external tensorrt_plugins library from EP directory (same location as nvinfer_plugin)
  // This library contains GroupQueryAttention and RotaryEmbedding plugins for transformer models
  try {
    const auto& env = onnxruntime::GetDefaultEnv();
    auto external_plugin_path = GetEPLibraryDirectory() +
                                PathString(LIBRARY_PREFIX ORT_TSTR("tensorrt_plugins") LIBRARY_EXTENSION);
    void* external_plugin_handle = nullptr;
    auto status = env.LoadDynamicLibrary(external_plugin_path, false, &external_plugin_handle);
    if (status.IsOK()) {
      LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] External plugins loaded: tensorrt_plugins (GQA + RotaryEmbedding)";
    } else {
      LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] tensorrt_plugins library not found in EP library path (optional)";
    }
  } catch (const std::exception& e) {
    LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] tensorrt_plugins library not available: " << e.what();
  }

  try {
    // Get all registered TRT plugins from registry
    LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] Getting all registered TRT plugins from TRT plugin registry ...";
    TensorrtLogger trt_logger = GetTensorrtLogger(false);
    void* library_handle = nullptr;
    const auto& env = onnxruntime::GetDefaultEnv();
    auto full_path = env.GetRuntimePath() +
                     PathString(LIBRARY_PREFIX ORT_TSTR("nvinfer_plugin") LIBRARY_EXTENSION);
    ORT_THROW_IF_ERROR(env.LoadDynamicLibrary(full_path, false, &library_handle));

    bool (*dyn_initLibNvInferPlugins)(void* logger, char const* libNamespace);
    ORT_THROW_IF_ERROR(env.GetSymbolFromLibrary(library_handle, "initLibNvInferPlugins", (void**)&dyn_initLibNvInferPlugins));
    dyn_initLibNvInferPlugins(&trt_logger, "");
    LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] Default plugins successfully loaded.";

#if defined(_MSC_VER)
#pragma warning(push)
#pragma warning(disable : 4996)  // Ignore warning C4996: 'nvinfer1::*' was declared deprecated
#endif
  } catch (const std::exception&) {
    LOGS_DEFAULT(INFO) << "[NvTensorRTRTX EP] Default plugin library is not on the path and is therefore ignored";
  }
  try {
    // getAllCreators() is the TRT 10+ replacement for the removed getPluginCreatorList().
    // It returns IPluginCreatorInterface* const*; cast each entry to the deprecated-but-present
    // IPluginCreator* to access getPluginName()/getPluginVersion().
    int32_t num_plugin_creator = 0;
    auto plugin_creators = getPluginRegistry()->getAllCreators(&num_plugin_creator);
    std::unordered_set<std::string> registered_plugin_names;

    for (int32_t i = 0; i < num_plugin_creator; i++) {
      auto* plugin_creator = dynamic_cast<nvinfer1::IPluginCreator*>(plugin_creators[i]);
      if (!plugin_creator) continue;  // skip IPluginCreatorV3One entries that lack these accessors
      std::string plugin_name(plugin_creator->getPluginName());
      LOGS_DEFAULT(VERBOSE) << "[NvTensorRTRTX EP] " << plugin_name << ", version : " << plugin_creator->getPluginVersion();

      // plugin has different versions and we only register once
      if (registered_plugin_names.find(plugin_name) != registered_plugin_names.end()) {
        continue;
      }

      created_custom_op_list.push_back(std::make_unique<TensorRTCustomOp>(onnxruntime::kNvTensorRTRTXExecutionProvider, nullptr));  // Make sure TensorRTCustomOp object won't be cleaned up
      created_custom_op_list.back().get()->SetName(plugin_creator->getPluginName());
      custom_op_domain->custom_ops_.push_back(created_custom_op_list.back().get());
      registered_plugin_names.insert(plugin_name);
    }

#if defined(_MSC_VER)
#pragma warning(pop)
#endif

    custom_op_domain->domain_ = "trt.plugins";
    domain_list.push_back(custom_op_domain.get());
  } catch (const std::exception&) {
    LOGS_DEFAULT(WARNING) << "[NvTensorRTRTX EP] Failed to get TRT plugins from TRT plugin registration. Therefore, TRT EP can't create custom ops for TRT plugins";
  }

  // Register native custom ops (register these independent of TRT plugin library availability)
  if (native_custom_op_list.empty()) {
    const char* native_custom_ops_names[] = {"TRT_FP4DynamicQuantize", "TRT_FP8QuantizeLinear", "TRT_FP8DequantizeLinear"};
    size_t num_native_custom_ops = std::size(native_custom_ops_names);

    for (size_t i = 0; i < num_native_custom_ops; i++) {
      native_custom_op_list.push_back(std::make_unique<TensorRTCustomOp>(onnxruntime::kNvTensorRTRTXExecutionProvider, nullptr));
      native_custom_op_list.back()->SetName(native_custom_ops_names[i]);
      native_custom_op_domain->custom_ops_.push_back(native_custom_op_list.back().get());
    }

    native_custom_op_domain->domain_ = "trt";
    domain_list.push_back(native_custom_op_domain.get());
  }

  return Status::OK();
}

void ReleaseTensorRTCustomOpDomain(OrtCustomOpDomain* domain) {
  (void)domain;  // Suppress unused parameter warning
  // The domain and its custom ops are owned by static unique_ptrs in CreateTensorRTCustomOpDomainList().
  // Callers receive raw pointers via .get().
  //  1. Manually deleting them would cause a double-free when the static unique_ptrs are destroyed at program exit.
  //  2. Resetting the static unique_ptrs is also unsafe because other EP instances or InferenceSession objects
  //     may still hold raw pointers to these same objects (handed out via domain_list).
  // The static objects would be shared across EP instances and would persist for the program lifetime.
}

void ReleaseTensorRTCustomOpDomainList(std::vector<OrtCustomOpDomain*>& custom_op_domain_list) {
  // Only clear the reference vector, don't delete the static domain objects.
  custom_op_domain_list.clear();
}

}  // namespace onnxruntime
