#include <maps/wikimap/mapspro/services/tasks_yt/src/diffalert_worker/lib/worker.h>

#include <maps/libs/cmdline/include/cmdline.h>
#include <maps/libs/log8/include/log8.h>
#include <maps/libs/mongo/include/init.h>
#include <maps/tools/grinder/worker/include/api.h>

#include <maps/libs/common/include/environment.h>
#include <maps/libs/common/include/exception.h>
#include <yandex/maps/wiki/common/default_config.h>
#include <yandex/maps/wiki/common/extended_xml_doc.h>
#include <yandex/maps/wiki/common/yt.h>
#include <yandex/maps/wiki/tasks/revocation_checker.h>
#include <yandex/maps/wiki/tasks/tasks.h>
#include <yandex/maps/wiki/tasks/yt.h>

#include <util/generic/size_literals.h>


namespace common = maps::wiki::common;
namespace configs = maps::wiki::configs;
namespace worker = maps::grinder::worker;
namespace tasks = maps::wiki::tasks;

namespace maps::wiki::diffalert_worker {

const std::string TASK_TYPE = "diffalert";
const size_t WORKER_CONCURRENCY = 1;

const auto EDITOR_CONFIG_XPATH = "/config/services/editor/config";
const auto POSTGRESQL_MAPSPRO_MDB = "POSTGRESQL_MAPSPRO_MDB";
const auto TASK_ID = "taskId";

const auto YT_CLUSTER = "hahn";
const auto YT_WORKER_POOL = "maps-nmaps-validation";
const auto YT_WORKER_MEMORY = 64_GB;
const auto YT_WORKER_THREADS = 15;
const auto YT_WORKER_DISK_SPACE_LIMIT = 32_GB;
const auto YT_ACL_GROUP = "maps-nmaps-pool";

namespace {
tasks::TaskStatus convertStatus(Worker::Status status)
{
    switch (status)
    {
        case Worker::Status::Ok:
            return tasks::TaskStatus::Success;
        case Worker::Status::Canceled:
            return tasks::TaskStatus::Revoked;
        case Worker::Status::Failed:
            return tasks::TaskStatus::Failed;
    }
}
} // anonymous namespace

void runLocal(
    const common::ExtendedXmlDoc& config,
    const configs::editor::ConfigHolder& editorConfig,
    const worker::Task& task)
{
    const auto& args = task.args();
    auto taskId = args[TASK_ID].as<DBID>();

    Worker worker(config, editorConfig);
    auto taskData = worker.loadTaskData(taskId);
    taskData.dump();

    auto status = worker.run(
        taskData, [&] { return task.isCanceled(); });
    ASSERT(tasks::setOverriddenStatus(worker.corePool(), taskId, convertStatus(status)));
    if (status == Worker::Status::Canceled) {
        throw worker::TaskCanceledException();
    }
    REQUIRE(status == Worker::Status::Ok, "Task " << taskId << " failed");
}

void runLocal(
    const common::ExtendedXmlDoc& config,
    const configs::editor::ConfigHolder& editorConfig,
    DBID taskId)
{
    Worker worker(config, editorConfig);
    auto taskData = worker.loadTaskData(taskId);
    taskData.dump();

    auto status = worker.run(taskData, [&] { return false; });
    ASSERT(tasks::setOverriddenStatus(worker.corePool(), taskId, convertStatus(status)));
    REQUIRE(status == Worker::Status::Ok, "Task " << taskId << " failed");
}


class TVanillaDiffalert : public NYT::IVanillaJob<>
{
public:
    TVanillaDiffalert() = default;
    TVanillaDiffalert(TStringBuf environment, DBID taskId)
        : environment_(environment)
        , taskId_(taskId)
    {}

    void Do() override
    {
        auto configPtr = tasks::yt::loadConfigFromEnvironment(environment_);
        configs::editor::ConfigHolder editorConfig;
        Worker worker(*configPtr, editorConfig);
        auto taskData = worker.loadTaskData(taskId_);
        taskData.dump();
        if (taskData.alreadyStarted) {
            WARN() <<  "Task " << taskId_ << " has been already started somewhere else";

            // fix up broken previous tasks, do nothing if final status is already set
            tasks::setOverriddenStatus(worker.corePool(), taskId_, tasks::TaskStatus::Failed);
            return;
        }

        tasks::RevocationChecker revocationChecker(worker.corePool(), taskId_);
        auto status = worker.run(taskData, revocationChecker.wrapIntoCallable());
        ASSERT(tasks::setOverriddenStatus(worker.corePool(), taskId_, convertStatus(status)));
        REQUIRE(status == Worker::Status::Ok, "Task " << taskId_ << " failed");
    }

    Y_SAVELOAD_JOB(environment_, taskId_);

private:
    TString environment_;
    DBID taskId_ = 0;
};

REGISTER_VANILLA_JOB(TVanillaDiffalert);


void launchLocal(
    const common::ExtendedXmlDoc& config,
    const configs::editor::ConfigHolder& editorConfig,
    const worker::Task& task)
{
    const auto& args = task.args();
    auto taskId = args[TASK_ID].as<DBID>();
    ASSERT(tasks::setOverriddenStatus(config, taskId, tasks::TaskStatus::InProgress));
    runLocal(config, editorConfig, task);
}

void launchLocal(
    const common::ExtendedXmlDoc& config,
    const configs::editor::ConfigHolder& editorConfig,
    DBID taskId)
{
    ASSERT(tasks::setOverriddenStatus(config, taskId, tasks::TaskStatus::InProgress));
    runLocal(config, editorConfig, taskId);
}

void launchInYt(
    const common::ExtendedXmlDoc& config,
    DBID taskId)
{
    TString environment(toString(maps::common::getYandexEnvironment()));
    INFO() << "Received task: " << taskId << " " << environment;

    auto corePoolHolder = tasks::corePoolHolder(config);
    auto&& corePool = corePoolHolder.pool();

    auto vault = NYT::TNode::CreateMap()
        (POSTGRESQL_MAPSPRO_MDB, common::getCoreDbPassword(config).c_str());

    auto client = common::yt::createYtClient(YT_CLUSTER);

    auto title = TASK_TYPE + " (" + environment + ", " + std::to_string(taskId) + ")";

    auto acl = NYT::TNode().Add(
        NYT::TNode()
            ("subjects", NYT::TNode().Add(YT_ACL_GROUP))
            ("action", "allow")
            ("permissions", NYT::TNode().Add("read"))
    );

    auto taskSpec = NYT::TUserJobSpec()
        .MemoryLimit(YT_WORKER_MEMORY)
        .CpuLimit(YT_WORKER_THREADS)
        .MemoryReserveFactor(1)
        .DiskSpaceLimit(YT_WORKER_DISK_SPACE_LIMIT);

    auto task = NYT::TVanillaTask()
        .JobCount(1)
        .Name(TASK_TYPE.c_str())
        .Spec(taskSpec)
        .Job(new TVanillaDiffalert(environment, taskId));

    auto operationSpec = NYT::TNode::CreateMap()
        ("pool", YT_WORKER_POOL)
        ("title", title.c_str())
        ("acl", acl);

    ASSERT(tasks::setOverriddenStatus(corePool, taskId, tasks::TaskStatus::InProgress));
    try {
        auto operationPtr = client->RunVanilla(
            NYT::TVanillaOperationSpec()
                .AddTask(task),
            NYT::TOperationOptions()
                .Wait(false)
                .Spec(operationSpec)
                .SecureVault(vault));

        ASSERT(operationPtr.Get());
        auto operationId = operationPtr->GetId();
        tasks::yt::saveTaskOperationId(corePool, taskId, operationId);
        INFO() << "Launched taskId: " << taskId << " YT: " << GetGuidAsString(operationId);
    } catch (const std::exception& ex) {
        ERROR() << "Launch failed, taskId: " << taskId << " error: " << ex.what();
        tasks::setOverriddenStatus(corePool, taskId, tasks::TaskStatus::Failed, true); //force
    }
}

void launchInYt(
    const common::ExtendedXmlDoc& config,
    const worker::Task& task)
{
    const auto& args = task.args();
    auto taskId = args[TASK_ID].as<DBID>();

    launchInYt(config, taskId);
}

} // namespace maps::wiki::diffalert_worker


int main(int argc, char** argv)
{
    using namespace maps::wiki::diffalert_worker;

    try {
        NYT::Initialize(argc, argv);

        maps::cmdline::Parser parser;
        auto syslogTag = parser.string("syslog-tag")
            .help("redirect log output to syslog with given tag");
        auto servicesConfig = parser.file("config")
            .help("path to services configuration");
        auto grinderConfig = parser.file("grinder-config")
            .help("path to grinder configuration");
        auto taskId = parser.size_t("task-id")
            .help("service task id");
        auto yt = parser.flag("yt").defaultValue(false)
            .help("launch yt task");
        parser.parse(argc, argv);

        if (syslogTag.defined()) {
            maps::log8::setBackend(maps::log8::toSyslog(syslogTag));
        }

        maps::mongo::init();

        std::unique_ptr<common::ExtendedXmlDoc> configDocPtr;
        if (servicesConfig.defined()) {
            configDocPtr = std::make_unique<common::ExtendedXmlDoc>(servicesConfig);
        } else {
            configDocPtr = common::loadDefaultConfig();
        }

        auto editorConfig = [&] {
            return configs::editor::ConfigHolder(
                configDocPtr->get<std::string>(EDITOR_CONFIG_XPATH));
        };

        if (taskId.defined()) {
            if (yt) {
                launchInYt(*configDocPtr, taskId);
            } else {
                launchLocal(*configDocPtr, editorConfig(), taskId);
            }
            return EXIT_SUCCESS;
        }

        worker::Options options;
        if (grinderConfig.defined()) {
            options.setConfigLocation(grinderConfig);
        }
        options.setConcurrencyLevel(WORKER_CONCURRENCY);

        options.on(
            TASK_TYPE,
            [&](const worker::Task& task) {
                if (yt) {
                    launchInYt(*configDocPtr, task);
                } else {
                    launchLocal(*configDocPtr, editorConfig(), task);
                }
            });

        worker::run(std::move(options));
    } catch (const maps::Exception& ex) {
        FATAL() << "Worker failed: " << ex;
        return EXIT_FAILURE;
    } catch (const std::exception& ex) {
        FATAL() << "Worker failed: " << ex.what();
        return EXIT_FAILURE;
    }
    return EXIT_SUCCESS;
}
