#include "nodes/track/trackNode.hpp"

namespace Node
{


void TrackNode::work()
{
    printf("TrackNode %s\n", name_.c_str());
    while (running_)
    {
        bool has_data = false;
        for (auto& input_buffer : input_buffers_)
        {
            std::shared_ptr<meta::MetaData> metaData;
            if (!input_buffer.second->try_pop(metaData))
            {
                continue;
            }
            has_data = true;
            // printf("Node %s get data from %s\n", name_.c_str(), input_buffer.first.c_str());

            // auto res = model_->forward(tensor::cvimg(image), image.cols, image.rows, 0.0f, 0.0f);
            if (!track_)
            {
                printf("track is nullptr\n");
                continue;
            }
            std::vector<Object> objects;
            std::transform(metaData->boxes.begin(), metaData->boxes.end(), objects.begin(), [](data::Box& box) {
                Object obj;
                obj.rect.x = box.left;
                obj.rect.y = box.top;
                obj.rect.width = box.right - box.left;
                obj.rect.height = box.bottom - box.top;
                obj.label = box.class_id;
                obj.prob = box.score;
                return obj;
            });
            std::vector<STrack> output_stracks = tracker_->update(objects);
            for (int i = 0; i < output_stracks.size(); i++)
            {
                vector<float> tlwh = output_stracks[i].tlwh;
                int track_id = output_stracks[i].track_id;
                std::for_each(metaData->boxes.begin(), metaData->boxes.end(), [track_id, tlwh](data::Box& box) {
                    if (box.left == tlwh[0] && box.top == tlwh[1] && box.right - box.left == tlwh[2] && box.bottom - box.top == tlwh[3])
                    {
                        box.track_id = track_id;
                    }
                });
            }
        }
        if (!has_data)
        {
            std::unique_lock<std::mutex> lock(mutex_);
            cond_var_->wait_for(lock, std::chrono::milliseconds(100), [this] {
                return !running_;  // 等待时检查退出条件
            });
        }
    }
};

}   // namespace Node