#pragma once

#include <NvInfer.h>
#include <NvInferVersion.h>
#include <cuda_runtime.h>

#include <cstring>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>

namespace fangjia {

inline void checkCuda(cudaError_t status, char const* expression, char const* file, int line) {
    if (status != cudaSuccess) {
        throw std::runtime_error(
            std::string("CUDA failure: ") + expression + " at " + file + ":" +
            std::to_string(line) + ": " + cudaGetErrorString(status));
    }
}

#define CUDA_CHECK(expr) ::fangjia::checkCuda((expr), #expr, __FILE__, __LINE__)

class TrtLogger final : public nvinfer1::ILogger {
public:
    explicit TrtLogger(Severity minimum = Severity::kWARNING) : minimum_(minimum) {}

    void log(Severity severity, char const* message) noexcept override {
        if (severity <= minimum_) {
            std::cerr << "[TensorRT] " << message << '\n';
        }
    }

private:
    Severity minimum_;
};

template <typename T>
struct TrtDeleter {
    void operator()(T* object) const noexcept { delete object; }
};

template <typename T>
using TrtUniquePtr = std::unique_ptr<T, TrtDeleter<T>>;

inline std::vector<std::uint8_t> readBinary(std::filesystem::path const& path) {
    std::ifstream stream(path, std::ios::binary | std::ios::ate);
    if (!stream) {
        throw std::runtime_error("Cannot open " + path.string());
    }
    auto const size = stream.tellg();
    if(size<=0)throw std::runtime_error("Empty or unreadable engine: "+path.string());
    stream.seekg(0);
    std::vector<std::uint8_t> bytes(static_cast<std::size_t>(size));
    if(!stream.read(reinterpret_cast<char*>(bytes.data()), size))throw std::runtime_error("Truncated engine: "+path.string());
    return bytes;
}

inline void writeBinary(std::filesystem::path const& path, void const* data, std::size_t size) {
    std::filesystem::create_directories(path.parent_path());
    std::ofstream stream(path, std::ios::binary);
    if (!stream) {
        throw std::runtime_error("Cannot write " + path.string());
    }
    stream.write(static_cast<char const*>(data), static_cast<std::streamsize>(size));
}

inline std::size_t volume(nvinfer1::Dims const& dims) {
    std::size_t value = 1;
    for (int index = 0; index < dims.nbDims; ++index) {
        if (dims.d[index] < 0) {
            throw std::runtime_error("Unresolved dynamic tensor dimension");
        }
        value *= static_cast<std::size_t>(dims.d[index]);
    }
    return value;
}

inline std::size_t dataTypeSize(nvinfer1::DataType type) {
    switch (type) {
    case nvinfer1::DataType::kFLOAT: return 4;
    case nvinfer1::DataType::kHALF: return 2;
    case nvinfer1::DataType::kINT8: return 1;
    case nvinfer1::DataType::kINT32: return 4;
#if NV_TENSORRT_MAJOR >= 10
    case nvinfer1::DataType::kINT64: return 8;
#endif
    case nvinfer1::DataType::kBOOL: return 1;
    default: throw std::runtime_error("Unsupported TensorRT data type");
    }
}

class CudaEventTimer {
public:
    CudaEventTimer() {
        CUDA_CHECK(cudaEventCreate(&start_));
        CUDA_CHECK(cudaEventCreate(&end_));
    }
    ~CudaEventTimer() {
        cudaEventDestroy(start_);
        cudaEventDestroy(end_);
    }
    void start(cudaStream_t stream) { CUDA_CHECK(cudaEventRecord(start_, stream)); }
    float stop(cudaStream_t stream) {
        CUDA_CHECK(cudaEventRecord(end_, stream));
        CUDA_CHECK(cudaEventSynchronize(end_));
        float milliseconds = 0.0F;
        CUDA_CHECK(cudaEventElapsedTime(&milliseconds, start_, end_));
        return milliseconds;
    }

private:
    cudaEvent_t start_{};
    cudaEvent_t end_{};
};

} // namespace fangjia