#pragma once

#include "common.hpp"

#include <unordered_map>

namespace fangjia {

class TrtEngine {
public:
    TrtEngine(std::filesystem::path const& enginePath, TrtLogger& logger);
    ~TrtEngine();

    TrtEngine(TrtEngine const&) = delete;
    TrtEngine& operator=(TrtEngine const&) = delete;

    void setInputShape(std::string const& name, nvinfer1::Dims const& shape);
    void prepare();
    void bindExternal(std::string const& name, void* pointer, std::size_t bytes);
    void useOwned(std::string const& name);
    bool enqueue(cudaStream_t stream);

    void* device(std::string const& name);
    std::size_t bytes(std::string const& name) const;
    nvinfer1::Dims shape(std::string const& name) const;
    nvinfer1::DataType dataType(std::string const& name) const;
    std::vector<std::string> inputs() const;
    std::vector<std::string> outputs() const;

private:
    struct Buffer {
        void* owned{};
        void* active{};
        std::size_t capacity{};
        std::size_t required{};
        bool external{};
    };

    Buffer& buffer(std::string const& name);
    Buffer const& buffer(std::string const& name) const;

    TrtUniquePtr<nvinfer1::IRuntime> runtime_;
    TrtUniquePtr<nvinfer1::ICudaEngine> engine_;
    TrtUniquePtr<nvinfer1::IExecutionContext> context_;
    std::filesystem::path enginePath_;
    std::unordered_map<std::string, Buffer> buffers_;
};

} // namespace fangjia