#include "trt_engine.hpp"

#include <algorithm>

namespace fangjia {

TrtEngine::TrtEngine(std::filesystem::path const& enginePath, TrtLogger& logger)
    : enginePath_(enginePath) {
    auto bytes = readBinary(enginePath);
    runtime_.reset(nvinfer1::createInferRuntime(logger));
    if (!runtime_) throw std::runtime_error("createInferRuntime failed");
    engine_.reset(runtime_->deserializeCudaEngine(bytes.data(), bytes.size()));
    if (!engine_) throw std::runtime_error("Cannot deserialize " + enginePath.string());
    context_.reset(engine_->createExecutionContext());
    if (!context_) throw std::runtime_error("createExecutionContext failed");
    for (int index = 0; index < engine_->getNbIOTensors(); ++index) {
        buffers_.emplace(engine_->getIOTensorName(index), Buffer{});
    }
}

TrtEngine::~TrtEngine() {
    for (auto& [name, item] : buffers_) {
        (void) name;
        if (item.owned) cudaFree(item.owned);
    }
}

void TrtEngine::setInputShape(std::string const& name, nvinfer1::Dims const& shape) {
    if (!context_->setInputShape(name.c_str(), shape)) {
        throw std::runtime_error("TensorRT rejected input shape for " + name);
    }
}

void TrtEngine::prepare() {
    for (int index = 0; index < engine_->getNbIOTensors(); ++index) {
        std::string const name = engine_->getIOTensorName(index);
        auto const dims = context_->getTensorShape(name.c_str());
        for (int axis = 0; axis < dims.nbDims; ++axis) {
            if (dims.d[axis] < 0) {
                throw std::runtime_error(
                    "TensorRT dimensions are incomplete for " + enginePath_.string() +
                    ":" + name);
            }
        }
        auto& item = buffer(name);
        item.required = volume(dims) * dataTypeSize(engine_->getTensorDataType(name.c_str()));
        if (item.capacity < item.required) {
            if (item.owned) CUDA_CHECK(cudaFree(item.owned));
            CUDA_CHECK(cudaMalloc(&item.owned, item.required));
            item.capacity = item.required;
        }
        item.active = item.owned;
    }
}

bool TrtEngine::enqueue(cudaStream_t stream) {
    for (int index = 0; index < engine_->getNbIOTensors(); ++index) {
        std::string const name = engine_->getIOTensorName(index);
        auto& item = buffer(name);
        if (!item.active) throw std::runtime_error("No device buffer for " + name);
        if (!context_->setTensorAddress(name.c_str(), item.active)) {
            throw std::runtime_error("setTensorAddress failed for " + name);
        }
    }
    return context_->enqueueV3(stream);
}

void* TrtEngine::device(std::string const& name) { return buffer(name).active; }
std::size_t TrtEngine::bytes(std::string const& name) const { return buffer(name).required; }
nvinfer1::Dims TrtEngine::shape(std::string const& name) const {
    return context_->getTensorShape(name.c_str());
}
nvinfer1::DataType TrtEngine::dataType(std::string const& name) const {
    return engine_->getTensorDataType(name.c_str());
}

std::vector<std::string> TrtEngine::inputs() const {
    std::vector<std::string> result;
    for (int index = 0; index < engine_->getNbIOTensors(); ++index) {
        char const* name = engine_->getIOTensorName(index);
        if (engine_->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT) result.emplace_back(name);
    }
    return result;
}

std::vector<std::string> TrtEngine::outputs() const {
    std::vector<std::string> result;
    for (int index = 0; index < engine_->getNbIOTensors(); ++index) {
        char const* name = engine_->getIOTensorName(index);
        if (engine_->getTensorIOMode(name) == nvinfer1::TensorIOMode::kOUTPUT) result.emplace_back(name);
    }
    return result;
}

TrtEngine::Buffer& TrtEngine::buffer(std::string const& name) {
    auto iterator = buffers_.find(name);
    if (iterator == buffers_.end()) throw std::runtime_error("Unknown tensor " + name);
    return iterator->second;
}
TrtEngine::Buffer const& TrtEngine::buffer(std::string const& name) const {
    auto iterator = buffers_.find(name);
    if (iterator == buffers_.end()) throw std::runtime_error("Unknown tensor " + name);
    return iterator->second;
}

} // namespace fangjia