#include <yandex/maps/wiki/threadutils/scheduler.h>
#include <yandex/maps/wiki/threadutils/threadedqueue.hpp>

#include <algorithm>
#include <list>
#include <unordered_set>
#include <exception>

namespace maps::wiki {

class Scheduler::Impl
{
public:
    Impl() = default;

    TTaskId addTask(
            const Runner& runner,
            const Executor& executor,
            const std::vector<TTaskId>& dependencies);

    void executeAll();

    void setIsCanceledChecker(CanceledChecker isCanceled);

private:
    struct Task
    {
        TTaskId id;
        Runner runner;
        Executor executor;
        std::vector<TTaskId> dependencies;
    };

    struct FinishedMessage
    {
        TTaskId id;
        std::exception_ptr error;
    };

private:
    void scheduleReady(
            const std::unordered_set<TTaskId>& finished,
            const std::shared_ptr<ThreadedQueue<FinishedMessage>>& finishedQueue);

private:
    TTaskId nextTaskId_{0};
    std::list<Task> pendingTasks_;
    CanceledChecker isCanceled_;
};


Scheduler::Scheduler()
    : impl_(new Impl)
{ }

Scheduler::~Scheduler() = default;

Scheduler::TTaskId Scheduler::addTask(
        const Runner& runner,
        const Executor& executor,
        const std::vector<TTaskId>& dependencies)
{
    return impl_->addTask(runner, executor, dependencies);
}

void Scheduler::executeAll()
{
    impl_->executeAll();
}

Scheduler& Scheduler::setIsCanceledChecker(CanceledChecker isCanceled)
{
    impl_->setIsCanceledChecker(std::move(isCanceled));
    return *this;
}

Scheduler::TTaskId Scheduler::Impl::addTask(
        const Runner& runner,
        const Executor& executor,
        const std::vector<TTaskId>& dependencies)
{
    Task task{nextTaskId_, runner, executor, dependencies};
    nextTaskId_++;
    pendingTasks_.push_back(task);
    return task.id;
}

void Scheduler::Impl::executeAll()
{
    size_t numTasks = pendingTasks_.size();
    if (numTasks == 0) {
        return;
    }

    std::shared_ptr<ThreadedQueue<FinishedMessage>> finishedQueue =
        std::make_shared<ThreadedQueue<FinishedMessage>>();
    ThreadedQueue<FinishedMessage>::RawContainer justFinished;
    std::unordered_set<TTaskId> finished;
    do {
        scheduleReady(finished, finishedQueue);
        finishedQueue->popAll(justFinished);
        for (const FinishedMessage& result : justFinished) {
            if (!result.error) {
                finished.insert(result.id);
            } else {
                std::rethrow_exception(result.error);
            }
        }
    } while (finished.size() < numTasks);
}

void Scheduler::Impl::setIsCanceledChecker(CanceledChecker isCanceled)
{
    isCanceled_ = std::move(isCanceled);
}

void Scheduler::Impl::scheduleReady(
        const std::unordered_set<Scheduler::TTaskId>& finished,
        const std::shared_ptr<ThreadedQueue<FinishedMessage>>& finishedQueue)
{
    for (auto taskIt = pendingTasks_.begin(); taskIt != pendingTasks_.end(); ) {
        if (std::all_of(
                taskIt->dependencies.begin(), taskIt->dependencies.end(),
                [&](TTaskId dep) { return finished.count(dep) > 0; })) {
            Task task = *taskIt;

            auto runner = [isCanceled = this->isCanceled_, finishedQueue, task]() {
                std::exception_ptr error;
                try {
                    if (isCanceled && isCanceled()) {
                        throw ExecutionCanceled();
                    }
                    task.runner();
                } catch (...) {
                    error = std::current_exception();
                }
                finishedQueue->push(FinishedMessage{task.id, error});
            };

            task.executor(runner);

            taskIt = pendingTasks_.erase(taskIt);
        } else {
            taskIt++;
        }
    }
}

} // namespace maps::wiki
