#include <iostream>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <vector>
#include <queue>
#include <functional>
#include <atomic>

class TaskScheduler {
private:
    std::vector<std::thread> workers;                    // Thread pool
    std::queue<std::function<void()>> taskQueue;         // Manual task queue
    std::mutex queueMutex;                               // Protects task queue
    std::condition_variable taskCV;                      // Notifies workers
    std::condition_variable completionCV;                // Notifies waitUntilComplete()
    std::atomic<int> activeTasks{0};                     // Number of active/scheduled tasks
    bool shutdown = false;                               // Flag for stopping workers

public:
    TaskScheduler(int numThreads = std::thread::hardware_concurrency()) {
        // Start a fixed number of worker threads
        for (int i = 0; i < numThreads; ++i) {
            workers.emplace_back([this]() {
                while (true) {
                    std::function<void()> task;

                    // Acquire task
                    {
                        std::unique_lock<std::mutex> lock(queueMutex);
                        taskCV.wait(lock, [this]() {
                            return shutdown || !taskQueue.empty();
                        });

                        if (shutdown && taskQueue.empty())
                            return;

                        task = std::move(taskQueue.front());
                        taskQueue.pop();
                    }

                    // Execute task outside lock
                    task();

                    // One task done
                    if (--activeTasks == 0) {
                        std::unique_lock<std::mutex> lock(queueMutex);
                        completionCV.notify_all();
                    }
                }
            });
        }
    }

    // Non-blocking function to schedule a task
    void schedule(std::function<void()> task) {
        {
            std::lock_guard<std::mutex> lock(queueMutex);
            taskQueue.push(std::move(task));
            ++activeTasks;
        }
        taskCV.notify_one(); // Notify one worker
    }

    // Blocking function until all tasks complete
    void waitUntilComplete() {
        std::unique_lock<std::mutex> lock(queueMutex);
        completionCV.wait(lock, [this]() {
            return activeTasks == 0;
        });
    }

    ~TaskScheduler() {
        {
            std::lock_guard<std::mutex> lock(queueMutex);
            shutdown = true;
        }
        taskCV.notify_all(); // Wake up all workers
        for (std::thread& t : workers)
            if (t.joinable()) t.join();
    }
};
int main() {
    TaskScheduler scheduler(4);

    for (int i = 0; i < 10; ++i) {
        scheduler.schedule([i]() {
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
            std::cout << "Task " << i << " completed.\n";
        });
    }

    scheduler.waitUntilComplete();
    std::cout << "All tasks completed.\n";

    return 0;
}