#ifndef RKYOLOV8_HPP
#define RKYOLOV8_HPP

#include <memory>
#include <opencv2/opencv.hpp>
#include <string>
#include <vector>

#include "rknn/rknn_api.h"
#include "rknn_test/postprocess.h"

// 引入我们刚才修复的 DMA 分配器
#include "rknn_test/dma_allocator.hpp"

namespace rknn_test {

class rkYolov8 {
public:
	rkYolov8(const std::string& model_path, const std::string& label_path, int class_num);
	~rkYolov8();

	int init(rknn_context* ctx_in, bool is_slave);
	rknn_context* get_pctx();
	detect_result_group_t infer(const cv::Mat& ori_img);

private:
	unsigned char* load_model(const char* filename, int* model_size);
	void post_process_v8_dfl(rknn_output* outputs, float scale, int pad_w, int pad_h,
							 detect_result_group_t* group);

private:
	std::string model_path;
	std::string m_label_path;
	int m_class_num;

	rknn_context ctx;
	bool is_slave = false;
	unsigned char* model_data = nullptr;

	rknn_input_output_num io_num;
	rknn_tensor_attr* input_attrs = nullptr;
	rknn_tensor_attr* output_attrs = nullptr;

	rknn_input inputs[1];

	int width = 0;
	int height = 0;
	int channel = 0;

	float conf_threshold;
	float nms_threshold;

	// 【修改点】使用 DmaBuffer 智能指针管理 NPU 输入内存
	std::unique_ptr<DmaBuffer> input_dma_buf_;
};

}  // namespace rknn_test

#endif	// RKYOLOV8_HPP