#include <thread>
#include <condition_variable>
#include <atomic>
#include <mutex>
#include <list>
#include <optional>

#include <Utilities/Async.hpp>
#include <CommandQueue/CommandQueue.hpp>
#include <Settings/Settings.hpp>
#include <Log/Log.hpp>
#include <Utilities/Utilities.hpp>
#include <Utilities/Perf.hpp>

typedef std::function<void()> Callback;

static std::atomic<int> numPendingTasks;
static int numWorkers;

class BarrierWaiter
{
private:
    std::atomic<int> workersLeft;
    Callback callback;

public:
    BarrierWaiter(const Callback &callback) : workersLeft(numWorkers), callback(callback) {}

    void notify()
    {
        if (--workersLeft == 0)
        {
            callback();
        }
    }
};

class Worker
{
private:
    static constexpr int maxNonMiscTasksInRow = 4;

    std::mutex mutex;
    std::condition_variable cond;
    std::list<Callback> miscQueue;
    std::vector<std::list<Callback>> chunkRenderQueues;
    std::vector<std::list<Callback>> worldGenQueues;
    int numNonMiscTasksInRow;
    std::atomic<int> ourPendingTasks;
    std::vector<std::vector<std::shared_ptr<BarrierWaiter>>> barrierWaiters;

    std::optional<Callback> tryPopNext(std::vector<std::list<Callback>> &queues)
    {
        size_t i;
        for (i=0; i<queues.size(); i++)
        {
            if (queues[i].size() != 0)
            {
                Callback cb = queues[i].front();
                queues[i].pop_front();
                return cb;
            }
        }

        return std::nullopt;
    }

    std::optional<Callback> tryPassBarrier(AsyncBarrier barrier)
    {
        std::vector<std::shared_ptr<BarrierWaiter>> &waiters = barrierWaiters[static_cast<int>(barrier)];
        if (waiters.size() == 0)
        {
            return std::nullopt;
        }

        Callback callback = [waiters = std::move(waiters)]()
        {
            for (std::shared_ptr<BarrierWaiter> waiter : waiters)
            {
                waiter->notify();
            }
        };

        waiters.clear();
        ourPendingTasks++;
        numPendingTasks++;
        return callback;
    }

    std::optional<Callback> getNextTask()
    {
        std::optional<Callback> cbBarrier = tryPassBarrier(AsyncBarrier::CurrentlyRunning);
        if (cbBarrier.has_value())
        {
            return cbBarrier;
        }
        
        if (miscQueue.size() == 0)
        {
            numNonMiscTasksInRow = 0;
        }

        if (numNonMiscTasksInRow >= maxNonMiscTasksInRow)
        {
            Callback cb = miscQueue.front();
            miscQueue.pop_front();
            numNonMiscTasksInRow = 0;
            return cb;
        }

        numNonMiscTasksInRow++;

        std::optional<Callback> result = tryPopNext(chunkRenderQueues);
        if (!result.has_value())
        {
            result = tryPopNext(worldGenQueues);
        }

        if (result.has_value())
        {
            return result;
        }

        if (miscQueue.size() == 0)
        {
            return tryPassBarrier(AsyncBarrier::CurrentlyDispatched);
        }

        numNonMiscTasksInRow = 0;
        Callback cb = miscQueue.front();
        miscQueue.pop_front();
        return cb;
    }

    void helperLoop()
    {
        for (;;)
        {
            std::unique_lock<std::mutex> lock(mutex);
            std::optional<Callback> nextTask;

            for (;;)
            {
                nextTask = getNextTask();
                if (!nextTask.has_value())
                {
                    cond.wait(lock);
                    continue;
                }
                else
                {
                    break;
                }
            }

            lock.unlock();
            (*nextTask)();
            numPendingTasks--;
            ourPendingTasks--;
        }
    }

public:
    Worker() :
        chunkRenderQueues(32),
        worldGenQueues(32),
        numNonMiscTasksInRow(0),
        barrierWaiters(static_cast<int>(AsyncBarrier::Num))
    {
        std::thread([this]()
        {
            this->helperLoop();
        }).detach();
    }

    void dispatch(std::function<void()> callback, float relativeDeadline, PriorityClass priorityClass)
    {
        std::unique_lock<std::mutex> lock(mutex);

        if (priorityClass == PriorityClass::Misc)
        {
            miscQueue.push_back(callback);
        }
        else if (priorityClass == PriorityClass::ChunkRender)
        {
            size_t queueIndex = std::min(chunkRenderQueues.size()-1, (size_t) (relativeDeadline * chunkRenderQueues.size()));
            chunkRenderQueues[queueIndex].push_back(callback);
        }
        else if (priorityClass == PriorityClass::WorldGen)
        {
            size_t queueIndex = std::min(worldGenQueues.size()-1, (size_t) (relativeDeadline * worldGenQueues.size()));
            worldGenQueues[queueIndex].push_back(callback);
        }
        else
        {
            crash("Invalid priority class!");
        }

        ourPendingTasks++;
        lock.unlock();
        cond.notify_one();
    }

    bool dispatchBarrierWaiter(AsyncBarrier barrier, std::shared_ptr<BarrierWaiter> waiter, bool forceAsync)
    {
        std::unique_lock<std::mutex> lock(mutex);

        if (!forceAsync && ourPendingTasks == 0)
        {
            waiter->notify();
            return false;
        }

        barrierWaiters[static_cast<int>(barrier)].push_back(waiter);
        lock.unlock();
        cond.notify_one();
        return true;
    }
};

static Worker *workers;
static std::atomic<int> nextTaskNum;

void initAsync()
{
    int numThreads;
    if (settings.numAsyncThreads.has_value())
    {
        numThreads = *settings.numAsyncThreads;
        Log::info(formatString("Using $+ async threads, based on the 'numAsyncThreads' setting.", numThreads));
    }
    else
    {
        unsigned int numHardwareThreads = std::thread::hardware_concurrency();
        if (numHardwareThreads == 0)
        {
            Log::error("Failed to detect the number of hardware threads; will use 2 async threads as fallback.");
            numThreads = 2;
        }
        else
        {
            numThreads = std::max(1u, numHardwareThreads - 1);
            Log::info(formatString("Using $+ async threads, based on hardware.", numThreads));
        }
    }

    workers = new Worker[numThreads];
    numWorkers = numThreads;
}

void dispatchTask(std::function<void()> callback, float relativeDeadline, PriorityClass priorityClass)
{
    PERF_EVENT(frame, runAsync,
    {
        numPendingTasks++;

        int workerIndex = (nextTaskNum++) % numWorkers;
        workers[workerIndex].dispatch(callback, relativeDeadline, priorityClass);
    })
}

void afterAsyncBarrier(AsyncBarrier barrier, std::function<void()> callback)
{
    std::shared_ptr<BarrierWaiter> waiter = std::make_shared<BarrierWaiter>(callback);

    int i;
    for (i=0; i<numWorkers; i++)
    {
        workers[i].dispatchBarrierWaiter(barrier, waiter, i == numWorkers - 1);
    }
}

void awaitAsyncBarrier(AsyncBarrier barrier)
{
    std::mutex mutex;
    std::condition_variable cv;
    bool done = false;

    afterAsyncBarrier(barrier, [&]() {
        std::lock_guard<std::mutex> lock(mutex);
        done = true;
        cv.notify_one();
    });

    std::unique_lock<std::mutex> lock(mutex);
    cv.wait(lock, [&]() { return done; });
}

int getPendingTasks()
{
    return numPendingTasks;
}