#pragma once

#include "pipeline_types.hpp"
#include "trt_engine.hpp"

#include <opencv2/core.hpp>

#include <string>

namespace fangjia {

class RtmwEngine {
public:
    RtmwEngine(
        std::filesystem::path const& enginePath,
        TrtLogger& logger,
        int maxSourceWidth,
        int maxSourceHeight,
        int maxPeople = 3);
    ~RtmwEngine();

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

    std::vector<PersonPose> infer(
        cv::Mat const& frame,
        std::vector<Box> const& boxes,
        cudaStream_t stream,
        float* preprocessMs = nullptr,
        float* inferenceMs = nullptr,
        float* postprocessMs = nullptr);

    void warmup(cudaStream_t stream);

private:
    TrtEngine engine_;
    int maxSourceWidth_{};
    int maxSourceHeight_{};
    int maxPeople_{};
    std::uint8_t* sourceDevice_{};
    std::uint8_t* cropsDevice_{};
    float* networkDevice_{};
    std::vector<double> centersHost_;
    std::vector<double> scalesHost_;
    std::vector<float> simccXHost_;
    std::vector<float> simccYHost_;
};

} // namespace fangjia