// Copyright chenyy 2026. Licensed under the MIT License.
// TensorRT logger header
//
// Licensed under the Apache License, Version 2.0 (the "License");

#ifndef ARMOR_YOLO_DETECT_TRT_LOGGER_HPP_
#define ARMOR_YOLO_DETECT_TRT_LOGGER_HPP_

#include <NvInfer.h>
#include <iostream>
#include <mutex>

class Logger : public nvinfer1::ILogger {
public:
  void log(Severity severity, const char* msg) noexcept override {
    std::lock_guard<std::mutex> lock(mutex_);
    switch (severity) {
      case Severity::kINTERNAL_ERROR:
        std::cerr << "[TensorRT] INTERNAL_ERROR: " << msg << std::endl;
        break;
      case Severity::kERROR:
        std::cerr << "[TensorRT] ERROR: " << msg << std::endl;
        break;
      case Severity::kWARNING:
        std::cerr << "[TensorRT] WARNING: " << msg << std::endl;
        break;
      case Severity::kINFO:
        std::cout << "[TensorRT] INFO: " << msg << std::endl;
        break;
      case Severity::kVERBOSE:
        std::cout << "[TensorRT] VERBOSE: " << msg << std::endl;
        break;
    }
  }

private:
  std::mutex mutex_;
};

extern Logger gLogger;

#endif  // ARMOR_YOLO_DETECT_TRT_LOGGER_HPP_
