#pragma once

#include "Program.hpp"
#include <thread>
#include <future>
#include <mutex>

namespace QuickMedia {
    template <class T, class... Args>
    class AsyncTask {
    public:
        using CallbackFunc = std::function<T(Args&&... args)>;

        AsyncTask() = default;

        AsyncTask(CallbackFunc callback_func, Args&&... args) {
            std::lock_guard<std::mutex> lock(mutex);
            std::promise<T> promise;
            future = promise.get_future();
            thread = std::thread(&AsyncTask::thread_handler, this, std::move(promise), std::move(callback_func), std::forward<Args>(args)...);
        }

        AsyncTask(AsyncTask &&other) noexcept {
            cancel();
            std::lock_guard<std::mutex> lock(mutex);
            thread = std::move(other.thread);
            future = std::move(other.future);
        }

        AsyncTask& operator=(AsyncTask &&other) noexcept {
            cancel();
            std::lock_guard<std::mutex> lock(mutex);
            thread = std::move(other.thread);
            future = std::move(other.future);
            return *this;
        }

        ~AsyncTask() {
            cancel();
        }

        bool valid() {
            std::lock_guard<std::mutex> lock(mutex);
            return future.valid();
        }

        bool ready() {
            std::lock_guard<std::mutex> lock(mutex);
            return future.valid() && future.wait_for(std::chrono::seconds(0)) == std::future_status::ready;
        }

        T get() {
            std::lock_guard<std::mutex> lock(mutex);
            if constexpr(std::is_same<T, void>::value) {
                if(thread.joinable()) {
                    thread.join();
                    future.get();
                }
            } else {
                T result = T();
                if(thread.joinable()) {
                    thread.join();
                    result = std::move(future.get());
                }
                return result;
            }
        }

        void cancel() {
            std::lock_guard<std::mutex> lock(mutex);
            if(future.valid()) {
                program_kill_in_thread(thread.get_id());
                if(thread.joinable()) {
                    thread.join();
                    future.get();
                }
            }
        }
    private:
        void thread_handler(std::promise<T> &&promise, CallbackFunc callback_func, Args&&... args) {
            if constexpr(std::is_same<T, void>::value) {
                callback_func(std::forward<Args>(args)...);
                promise.set_value();
            } else {
                promise.set_value(callback_func(std::forward<Args>(args)...));
            }
        }
    private:
        std::thread thread;
        std::future<T> future;
        std::mutex mutex;
    };
}