#include <NvInfer.h>
#include <NvInferPlugin.h>
#include <cuda_runtime.h>
#include <cmath>
#include <fstream>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>

struct Logger : nvinfer1::ILogger {
    void log(Severity level, const char* message) noexcept override {
        if (level <= Severity::kWARNING) std::cerr << message << '\n';
    }
};
void require(bool ok, const char* message) { if (!ok) throw std::runtime_error(message); }
void check(cudaError_t status) { if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); }
void shape(nvinfer1::Dims actual, const std::vector<int>& expected) {
    require(actual.nbDims == static_cast<int>(expected.size()), "rank mismatch");
    for (int i=0; i<actual.nbDims; ++i) require(actual.d[i] == expected[i], "shape mismatch");
}
struct Buffers {
    cudaStream_t stream=nullptr;
    std::vector<void*> values;
    ~Buffers() {
        if (stream) cudaStreamSynchronize(stream);
        for (void* ptr: values) if (ptr) cudaFree(ptr);
        if (stream) cudaStreamDestroy(stream);
    }
};
int main(int argc, char** argv) {
    try {
        require(argc==2, "expected engine path");
        Logger log; initLibNvInferPlugins(&log, "");
        std::ifstream file(argv[1], std::ios::binary);
        std::vector<char> bytes((std::istreambuf_iterator<char>(file)), {});
        require(!bytes.empty(), "empty engine");
        std::unique_ptr<nvinfer1::IRuntime> runtime(nvinfer1::createInferRuntime(log));
        require(bool(runtime), "runtime failure");
        std::unique_ptr<nvinfer1::ICudaEngine> engine(runtime->deserializeCudaEngine(bytes.data(), bytes.size()));
        require(bool(engine), "deserialize failure");
        require(engine->getNbIOTensors()==3 && engine->getNbOptimizationProfiles()==1, "tensor/profile count");
        const std::vector<std::string> names={"input", "dets", "labels"};
        const std::vector<std::vector<int>> specs={{-1,3,512,512},{-1,300,4},{-1,300,3}};
        for (int i=0;i<3;++i) {
            const char* name=names[i].c_str();
            shape(engine->getTensorShape(name),specs[i]);
            require(engine->getTensorDataType(name)==(i==0 ? nvinfer1::DataType::kHALF : nvinfer1::DataType::kFLOAT),"dtype mismatch");
            require(engine->getTensorIOMode(name)==(i==0 ? nvinfer1::TensorIOMode::kINPUT : nvinfer1::TensorIOMode::kOUTPUT),"mode mismatch");
        }
        shape(engine->getProfileShape("input",0,nvinfer1::OptProfileSelector::kMIN),{1,3,512,512});
        shape(engine->getProfileShape("input",0,nvinfer1::OptProfileSelector::kOPT),{4,3,512,512});
        shape(engine->getProfileShape("input",0,nvinfer1::OptProfileSelector::kMAX),{4,3,512,512});
        std::unique_ptr<nvinfer1::IExecutionContext> ctx(engine->createExecutionContext());
        require(bool(ctx),"context failure");
        for (int batch=1; batch<=4; ++batch) {
            require(ctx->setInputShape("input",nvinfer1::Dims4(batch,3,512,512)),"set shape failed");
            Buffers resources;check(cudaStreamCreate(&resources.stream));
            const std::vector<size_t> count={size_t(batch)*3*512*512,size_t(batch)*300*4,size_t(batch)*300*3};
            for (int i=0;i<3;++i) {
                auto expected=specs[i];expected[0]=batch;shape(ctx->getTensorShape(names[i].c_str()),expected);
                void* ptr=nullptr;check(cudaMalloc(&ptr,count[i]*(i==0 ? 2 : 4)));resources.values.push_back(ptr);
                require(ctx->setTensorAddress(names[i].c_str(),ptr),"set address failed");
            }
            check(cudaMemsetAsync(resources.values[0],0,count[0]*2,resources.stream));
            require(ctx->enqueueV3(resources.stream),"inference failed");check(cudaStreamSynchronize(resources.stream));
            for (int i=1;i<3;++i) {
                std::vector<float> output(count[i]);
                check(cudaMemcpy(output.data(),resources.values[i],count[i]*4,cudaMemcpyDeviceToHost));
                for (float v:output) require(std::isfinite(v),"nonfinite output");
            }
            std::cout << "{\"batch\":" << batch << ",\"shape_dtype_finite\":\"PASS\"}\n";
        }
        std::cout << "{\"contract\":\"PASS\",\"tensorrt\":\"" << NV_TENSORRT_MAJOR << '.' << NV_TENSORRT_MINOR
                  << '.' << NV_TENSORRT_PATCH << "\",\"accuracy\":\"NOT_RUN\"}\n";
        return 0;
    } catch (const std::exception& e) { std::cerr << e.what() << '\n'; return 1; }
}