#pragma once

#include "ofMain.h"
#include "ofxOnnxRuntime.h"

namespace ofxOnnxRuntime {
    class OnnxThread : public ofThread
    {
        public:
            ofxOnnxRuntime::BaseHandler* onnx;
            float* result = nullptr;
            bool isInferenceComplete = false;
            bool shouldRunInference = true;

            ~OnnxThread() {
                stop();
                waitForThread(false);
            }

            void setup(ofxOnnxRuntime::BaseHandler* onnx) {
                std::lock_guard<std::mutex> lock(mutex);
                this->onnx = onnx;
            }

            void start() {
                startThread();
            }

            void stop() {
                stopThread();
                condition.notify_all();
            }

            void threadedFunction() {
                while (isThreadRunning()) {
                    std::unique_lock<std::mutex> lock(mutex);
                    runOnnx();
                    condition.wait(lock);
                }
            }

            void update() {
                std::lock_guard<std::mutex> lock(mutex);
                condition.notify_one();
            }

            void runOnnx() {
                if (shouldRunInference) {
                    result = onnx->run();
                    isInferenceComplete = true;
                    shouldRunInference = false;
                }  
            }

            // Method to safely get the result
            float* getResult() {
                std::lock_guard<std::mutex> lock(mutex);
                return result;
            }

            bool checkInferenceComplete() {
                std::lock_guard<std::mutex> lock(mutex);
                return isInferenceComplete;
            }

            void resetInferenceFlag() {
                std::lock_guard<std::mutex> lock(mutex);
                isInferenceComplete = false;
            }

        protected:
            std::condition_variable condition;
    };
}