#include <QDebug>
#include <QThread>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <iomanip>
#include "aiengineinferenceopi5.h"
#include "file_utils.h"
#include "image_drawing.h"


AiEngineInferenceOpi5::AiEngineInferenceOpi5(QString modelPath, QObject *parent)
    : AiEngineInference{modelPath, parent}
{
    qDebug() << "AiEngineInferenceOpi5() test mModelPath=" << mModelPath;
}


void AiEngineInferenceOpi5::initialize(int number)
{
    mNumber = number;

    memset(&mRrknnAppCtx0, 0, sizeof(rknn_app_context_t));
    init_post_process();

    int ret = init_yolov8_model(mModelPath.toLocal8Bit(), &mRrknnAppCtx0);
    if (ret != 0) {
        qDebug() << "init_yolov8_model() failure! ret: " << ret << "modelPath = " << mModelPath << "number:" << number;
        return;
    }
}


AiEngineInferenceOpi5::~AiEngineInferenceOpi5()
{
    deinit_post_process();
    release_yolov8_model(&mRrknnAppCtx0);
}


image_buffer_t AiEngineInferenceOpi5::convertCV2FrameToImageBuffer(const cv::Mat& bgrFrame)
{
    // Convert BGR to RGB
    cv::Mat rgbFrame;
    cv::cvtColor(bgrFrame, rgbFrame, cv::COLOR_BGR2RGB);

    image_buffer_t imgBuffer;
    memset(&imgBuffer, 0, sizeof(image_buffer_t));
    imgBuffer.width         = rgbFrame.cols;
    imgBuffer.height        = rgbFrame.rows;
    imgBuffer.width_stride  = rgbFrame.step;
    imgBuffer.height_stride = rgbFrame.rows;
    imgBuffer.format        = IMAGE_FORMAT_RGB888;
    imgBuffer.size          = rgbFrame.total() * rgbFrame.elemSize();

    imgBuffer.virt_addr = new unsigned char[imgBuffer.size];
    std::memcpy(imgBuffer.virt_addr, rgbFrame.data, imgBuffer.size);

    return imgBuffer;
}


void AiEngineInferenceOpi5::freeImageBuffer(image_buffer_t& imgBuffer)
{
    if (imgBuffer.virt_addr) {
        delete[] imgBuffer.virt_addr;
        imgBuffer.virt_addr = nullptr;
    }
}


cv::Mat AiEngineInferenceOpi5::resizeToHalfAndAssigntoTopLeft640x640(const cv::Mat& inputFrame)
{
    // Resize input frame to half size
    cv::Mat resizedFrame;
    cv::resize(inputFrame, resizedFrame, cv::Size(), 0.5, 0.5);

    // Create a 640x640 frame to place the resized frame
    cv::Mat outputFrame = cv::Mat::zeros(640, 640, inputFrame.type());

    // Copy the resized frame to the top-left corner of the output frame
    cv::Rect roi(0, 0, resizedFrame.cols, resizedFrame.rows);
    resizedFrame.copyTo(outputFrame(roi));

    return outputFrame;
}


void AiEngineInferenceOpi5::drawObjects(cv::Mat& image, const object_detect_result_list& result_list)
{
    for (int i = 0; i < result_list.count; i++) {
        const object_detect_result& result = result_list.results[i];

        int left = result.box.left;
        int top = result.box.top;
        int right = result.box.right;
        int bottom = result.box.bottom;

        cv::rectangle(image, cv::Point(left, top), cv::Point(right, bottom), cv::Scalar(255, 0, 0), 2);

        // Text
        char c_text[256];
        sprintf(c_text, "%s %.1f%%", coco_cls_to_name(result.cls_id), result.prop * 100);
        cv::Point textOrg(left, top - 5);
        cv::putText(image, std::string(c_text), textOrg, cv::FONT_HERSHEY_SIMPLEX, result.prop, cv::Scalar(0, 0, 255), 1);
    }
}


void AiEngineInferenceOpi5::performInferenceSlot(cv::Mat frame)
{
    //qDebug() << "performInferenceSlot() in thread: " << QThread::currentThreadId();

    mActive = true;

    cv::Mat scaledFrame      = resizeToHalfAndAssigntoTopLeft640x640(frame);
    image_buffer_t imgBuffer = convertCV2FrameToImageBuffer(scaledFrame);

    object_detect_result_list od_results;
    int ret = inference_yolov8_model(&mRrknnAppCtx0, &imgBuffer, &od_results, mNumber);
    if (ret != 0) {
        qDebug() << "AiEngineInferenceOpi5::performInferenceSlot() failure! ret: " << ret;
        mActive = false;
        return;
    }

    AiEngineInferenceResult result;
    for (int i = 0; i < od_results.count; i++) {
        object_detect_result *det_result = &(od_results.results[i]);

        AiEngineObject object;
        object.classId          = det_result->cls_id;
        object.propability      = det_result->prop;
        object.rectangle.top    = det_result->box.top;
        object.rectangle.left   = det_result->box.left;
        object.rectangle.bottom = det_result->box.bottom;
        object.rectangle.right  = det_result->box.right;
        result.objects.append(object);
    }

    drawObjects(scaledFrame, od_results);
    freeImageBuffer(imgBuffer);

    result.frame = scaledFrame.clone();
    emit resultsReady(result);

    mActive = false;
}