#ifndef _YOLOV5_H
#define _YOLOV5_H

#include <iostream>
#include <fstream>
#include "utils.h"
#include "kpu.h"
#include "ai2d.h"

/**
 * @brief YOLOv5 Inference Class
 *
 * This class encapsulates the complete YOLOv5 detection process for each frame,
 * including preprocessing, model inference, and postprocessing to obtain and
 * visualize detection results.
 */
class Yolov5
{
public:
    /**
     * @brief Constructor
     *
     * Loads the YOLOv5 KModel, initializes the input/output tensors,
     * and sets up detection parameters.
     *
     * @param task_type      The type of task (e.g., "detect", "segment", "pose").
     * @param task_mode      The operation mode (e.g., "video", "image").
     * @param kmodel_file    Path to the KModel file.
     * @param image_wh       Original image width and height.
     * @param labels         Class label list.
     * @param conf_thres     Confidence threshold for detection filtering.
     * @param nms_thres      IoU threshold for Non-Maximum Suppression (NMS).
     * @param mask_thres     Mask threshold for segmentation tasks.
     * @param debug_mode     Debug level: 0 (off), 1 (timing only), 2 (verbose logs).
     */
    Yolov5(std::string task_type,
           std::string task_mode,
           std::string kmodel_file,
           FrameSize image_wh,
           std::vector<std::string> labels,
           float conf_thres = 0.5f,
           float nms_thres = 0.5f,
           float mask_thres = 0.5f,
           int debug_mode = 0);

    /**
     * @brief Destructor
     *
     * Releases all allocated resources.
     */
    ~Yolov5();

    /**
     * @brief Preprocess the input image
     *
     * Converts and resizes the image according to the model’s input format.
     * The processed result is stored in the model input tensor.
     *
     * @param input_tensor  Reference to the model input tensor.
     */
    void pre_process(runtime_tensor &input_tensor);

    /**
     * @brief Run model inference
     *
     * Executes the YOLOv5 KModel using the preprocessed input tensor.
     */
    void inference();

    /**
     * @brief Postprocess the model outputs
     *
     * Decodes raw inference outputs, applies confidence and NMS filtering,
     * and generates bounding boxes with class labels and confidence scores.
     *
     * @param yolo_results  Output vector containing detected bounding boxes.
     */
    void post_process(std::vector<YOLOBbox> &yolo_results);

    /**
     * @brief Draw detection results on the frame
     *
     * Renders bounding boxes, labels, and confidence values on the given image.
     *
     * @param draw_frame    Frame to draw on.
     * @param yolo_results  Detected bounding boxes from post-processing.
     */
    void draw_results(cv::Mat &draw_frame, std::vector<YOLOBbox> &yolo_results);

    /**
     * @brief Run complete detection pipeline
     *
     * Executes preprocessing, inference, and postprocessing in sequence.
     *
     * @param input_tensor  Input image tensor.
     * @param yolo_results  Output vector for detected bounding boxes.
     */
    void run(runtime_tensor &input_tensor,std::vector<YOLOBbox> &yolo_results);

private:
    runtime_tensor model_input_tensor_;   ///< AI2D output tensor used as model input
    std::string task_type_;               ///< Task type (e.g.,"classify", "detect", "segment")
    std::string task_mode_;               ///< Task mode (e.g., "video", "image")
    std::vector<std::string> labels_;     ///< Class label list
    int label_num_;                       ///< Number of class labels
    std::vector<cv::Scalar> colors;       ///< Colors used for drawing results
    FrameSize image_wh_;                  ///< Original image size
    FrameSize input_wh_;                  ///< Model input size
    float conf_thres_;                    ///< Confidence threshold
    float nms_thres_;                     ///< Non-Maximum Suppression (NMS) threshold
    float mask_thres_;                    ///< Mask threshold for segmentation
    int box_num_;                         ///< Number of detected boxes
    int max_box_num_;                     ///< Maximum allowed number of boxes
    int box_feature_len_;                 ///< Feature length per detection box
    float *output_det_;                   ///< Pointer to raw model output data
    int debug_mode_;                      ///< Debug mode level

    KPU *kpu = nullptr;                   ///< Pointer to KPU (Kendryte Processing Unit) instance
    AI2D *ai2d_yolo = nullptr;            ///< Pointer to AI2D preprocessor instance
};

#endif // _YOLOV5_H
