#pragma once

#include "params.h"
#include "categories.h"
#include "json_chunk.h"
#include "common.h"

#include <yandex/maps/wiki/threadutils/threadpool.h>
#include <yandex/maps/wiki/threadutils/threadedqueue.hpp>
#include <maps/libs/pgpool/include/pgpool3.h>

#include <boost/filesystem.hpp>

namespace mw = maps::wiki;
namespace bf = boost::filesystem;

namespace maps {

template <class T>
struct TaskWrapper {
    std::shared_ptr<T> task;

    void operator()()
    {
        task->run();
    }

    template <typename... Args>
    static TaskWrapper make(Args&&... args)
    {
        return TaskWrapper(std::make_shared<T>(std::forward<Args>(args)...));
    }

private:
    explicit TaskWrapper(std::shared_ptr<T> task)
        : task(std::move(task))
    {}
};

template <class T, typename... Args>
TaskWrapper<T> makeTask(Args&&... args)
{
    return TaskWrapper<T>::make(std::forward<Args>(args)...);
}

class Writer;
class JsonFileAttributes;

struct TaskContext {
    const Params& params;
    pgpool3::Pool& pgPool;
    mw::ThreadPool& loadIdsPool;
    mw::ThreadPool& loadRowsPool;
    mw::ThreadPool& fileWritePool;
    std::unique_ptr<Writer> writer;
    const JsonFileAttributes& jsonFileAttributes;
    bool failed;

    void onTaskError();

    void logPoolStats() const;
};

class TaskBase {
public:
    virtual ~TaskBase();
    void run();

protected:
    TaskBase(TaskContext& context, const std::string& name)
        : context_(context), name_(name)
    { }

private:
    virtual void doRun() = 0;

protected:
    TaskContext& context_;
    std::string name_;
};

class LoadIdsTask: public TaskBase {
public:
    LoadIdsTask(
            TaskContext& context,
            const Category& category)
        : TaskBase(context, "LoadIdsTask")
        , category_(category)
    { }

private:
    virtual void doRun();
    void runBatch(std::vector<DBID>& ids) const;

    const Category& category_;
};

class LoadRowsTask: public TaskBase {
public:
    LoadRowsTask(
            TaskContext& context,
            const Category& category,
            std::vector<DBID> ids)
        : TaskBase(context, "LoadRowsTask")
        , category_(category)
        , ids_(std::move(ids))
    { }

private:
    virtual void doRun();

    const Category& category_;
    std::vector<DBID>ids_;
};

class WriteFileTask: public TaskBase {
public:
    WriteFileTask(
            TaskContext& context,
            const bf::path& filePath,
            std::vector<JsonObject> objects)
        : TaskBase(context, "WriteFileTask")
        , filePath_(filePath)
        , objects_(std::move(objects))
    { }

private:
    virtual void doRun();

    bf::path filePath_;
    std::vector<JsonObject> objects_;
};

class WriteSingleFileTask: public TaskBase {
public:
    WriteSingleFileTask(
            TaskContext& context,
            mw::ThreadedQueue<JsonChunk>& queue)
        : TaskBase(context, "WriteSingleFileTask")
        , queue_(queue)
    { }

private:
    virtual void doRun();

    mw::ThreadedQueue<JsonChunk>& queue_;
};

class RegisterIdsTask: public TaskBase {
public:
    RegisterIdsTask(TaskContext& context)
        : TaskBase(context, "RegisterIdsTask")
    { }

private:
    virtual void doRun();
};

} // namespace maps
