#include "DetectAlgorithm.h"
#include "Core/Logger.h"
#include "ImageUtilities.h"
#include <ds_pedestrian_mot_hisi.h>
#include <filesystem>
#include <thread>

class DetectAlgorithmPrivate {
public:
    std::vector<io_TrackData> tracks;
};

DetectAlgorithm::DetectAlgorithm() : m_d(new DetectAlgorithmPrivate()) {
}

DetectAlgorithm::~DetectAlgorithm() {
    if (m_handle != nullptr) {
        ds_pedestrian_hisi_release(&m_handle);
    }
    if (m_d != nullptr) {
        delete m_d;
    }
}

DetectAlgorithm::Result DetectAlgorithm::detect(const uint8_t *nv21ImageData, uint64_t frameIndex) {
    DetectAlgorithm::Result ret;
    if (m_handle == nullptr) { // 一定得在这里执行
        initialize();
    }
    ImageUtilities::NV21ToBGR24(nv21ImageData, m_rgbImageBuffer.data(), DetectWidth, DetectHeight);

    std::vector<PedestrianRect> pedestrians;

    ds_pedestrian_det_hisi(m_handle, m_rgbImageBuffer.data(), pedestrians);
    ds_pedestrian_track_hisi(m_handle, m_rgbImageBuffer.data(), frameIndex, pedestrians, m_d->tracks);
    LOG(info) << "frame: " << frameIndex << ", pedestrians: " << pedestrians.size()
              << ", tracks: " << m_d->tracks.size();
    for (auto iterator = m_d->tracks.cbegin(); iterator != m_d->tracks.cend();) {
        if (iterator->track_state == TrackState::Remove) {
            ret.leaveTrackers.push_back(iterator->track_id);
            iterator = m_d->tracks.erase(iterator);
            continue;
        } else if (iterator->track_state == TrackState::Confirmed) {
            LOG(info) << iterator->prediction.body.conf << " " << iterator->prediction.head.conf;
            LOG(info) << iterator->prediction.body.state[0] << " " << iterator->prediction.body.state[1] << " "
                      << iterator->prediction.body.state[2] << " " << iterator->prediction.body.state[3];
            LOG(info) << iterator->prediction.head.state[0] << " " << iterator->prediction.head.state[1] << " "
                      << iterator->prediction.head.state[2] << " " << iterator->prediction.head.state[3];
        }
        ++iterator;
    }
    return ret;
}

void DetectAlgorithm::initialize() {
    constexpr auto licensePath = "/kdata/net.lic";
    bool licenseExisted = std::filesystem::exists(licensePath);
    if (licenseExisted && std::filesystem::file_size(licensePath) <= 0) {
        LOG(warning) << "license " << licensePath << " content is empty, remove it.";
        std::filesystem::remove(licensePath);
        licenseExisted = false;
    }
    ds_pedestrian_hisi_set_lic_path("/kdata");
    int status = ds_pedestrian_hisi_init(&m_handle, "/system/models/ds_mot_m0_2000.bin",
                                         "/system/models/ds_mot_m1_2000.bin", DetectWidth, DetectHeight, 3);
    if (status != 0) {
        LOG(error) << "ds_pedestrian_hisi_init() failed, status: " << status;
        m_handle = nullptr;
    } else {
        LOG(info) << "detect algorithm initialization successfully.";
    }
    if (!licenseExisted && std::filesystem::exists(licensePath)) {
        system("sync");
    }
}
