#include <maps/wikimap/mapspro/services/tasks_yt/src/validation_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 <yandex/maps/wiki/common/default_config.h>
#include <yandex/maps/wiki/common/extended_xml_doc.h>
#include <yandex/maps/wiki/common/pgpool3_helpers.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 validation = maps::wiki::validation;
namespace validator = maps::wiki::validator;
namespace worker = maps::grinder::worker;

namespace {

const std::string TASK_TYPE = "validation";
const std::string TASK_TYPE_HEAVY = "validation.heavy";

constexpr size_t WORKER_CONCURRENCY = 5;
constexpr size_t HEAVY_LOCAL_WORKER_CONCURRENCY = 1;

constexpr size_t YT_MAX_JOB_START_COUNT = 2;

const std::string& taskType(bool isHeavy)
{
    return isHeavy ? TASK_TYPE_HEAVY : TASK_TYPE;
}

} // anonymous namespace

namespace maps::wiki::validation {

const std::string EDITOR_CONFIG_XPATH = "/config/services/editor/config";
const std::string TASK_ID = "taskId";
const TString POSTGRESQL_MAPSPRO_MDB = "POSTGRESQL_MAPSPRO_MDB";

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

namespace {

tasks::TaskStatus convertStatus(Worker::Status workerStatus)
{
    switch (workerStatus)
    {
        case Worker::Status::Ok:
            return tasks::TaskStatus::Success;
        case Worker::Status::Canceled:
            return tasks::TaskStatus::Revoked;
        case Worker::Status::NeedRetry:
        case Worker::Status::Failed:
            return tasks::TaskStatus::Failed;
    }
}

} // anonymous namespace

void launchLocalWithSplit(
    const common::ExtendedXmlDoc& config,
    const validator::ValidatorConfig& validatorConfig,
    const worker::Task& task,
    bool isHeavy)
{
    const auto& args = task.args();
    auto taskId = args[TASK_ID].as<DBID>();

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

    worker.prepareTaskData(taskData);
    if (worker.splitTasks(taskData)) {
        return;
    }

    auto status = worker.run(taskData, [&] { return task.isCanceled(); });
    if (status == Worker::Status::Canceled) {
        throw worker::TaskCanceledException();
    }
}

void launchLocalWithoutSplit(
    const common::ExtendedXmlDoc& config,
    const validator::ValidatorConfig& validatorConfig,
    DBID taskId,
    bool isHeavy)
{
    Worker worker(config, validatorConfig, isHeavy);
    auto taskData = worker.loadTaskData(taskId);
    taskData.dump();

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

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

    void Do() override
    {
        auto configPtr = tasks::yt::loadConfigFromEnvironment(environment_);
        maps::wiki::validator::ValidatorConfig validatorConfig;

        Worker worker(*configPtr, validatorConfig, isHeavy_);
        auto taskData = worker.loadTaskData(taskId_);
        taskData.dump();
        if (taskData.startCount >= YT_MAX_JOB_START_COUNT) {
            WARN() << "Task " << taskId_ << " has failed too many times";
            tasks::setOverriddenStatus(worker.corePool(), taskId_, tasks::TaskStatus::Failed);
            return;
        }

        if (taskData.startCount > 0) {
            auto previousTaskStatus = tasks::getOverriddenStatus(worker.corePool(), taskId_);
            ASSERT(previousTaskStatus.has_value());
            if (*previousTaskStatus != tasks::TaskStatus::InProgress) {
                return;
            }

            worker.cleanResults(taskId_);
        }

        worker.setMaxCheckThreads(YT_WORKER_THREADS);
        worker.prepareTaskData(taskData);

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

    Y_SAVELOAD_JOB(environment_, taskId_, isHeavy_);

private:
    TString environment_;
    DBID taskId_ = 0;
    bool isHeavy_ = false;
};

REGISTER_VANILLA_JOB(TVanillaValidation);


void launchInYtWithoutSplit(
    const common::ExtendedXmlDoc& config,
    DBID taskId,
    bool isHeavy)
{
    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);

    const auto& taskType = ::taskType(isHeavy);
    auto title = taskType + " (" + 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(taskType.c_str())
        .Spec(taskSpec)
        .Job(new TVanillaValidation(environment, taskId, isHeavy));

    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 launchInYtWithSplit(
    const common::ExtendedXmlDoc& config,
    const validator::ValidatorConfig& validatorConfig,
    const worker::Task& task,
    bool isHeavy)
{
    const auto& args = task.args();
    auto taskId = args[TASK_ID].as<DBID>();
    {
        Worker worker(config, validatorConfig, isHeavy);
        try {
            auto taskData = worker.loadTaskData(taskId);
            taskData.dump();

            worker.prepareTaskData(taskData);
            if (worker.splitTasks(taskData)) {
                return;
            }
        } catch (const std::exception& ex) {
            ERROR() << "Launch failed, taskId: " << taskId << " error: " << ex.what();
            tasks::setOverriddenStatus(
                worker.corePool(), taskId, tasks::TaskStatus::Failed, true); //force
            return;
        }
    }
    launchInYtWithoutSplit(config, taskId, isHeavy);
}

} // namespace maps::wiki::validation


int main(int argc, char* argv[]) try
{
    NYT::Initialize(argc, argv);

    maps::cmdline::Parser parser;
    auto workerConfig = parser.string("config")
        .help("path to worker configuration");
    auto syslogTag = parser.string("syslog-tag")
        .help("redirect log output to syslog with given tag");
    auto grinderConfig = parser.string("grinder-config")
        .help("path to grinder configuration file");
    auto isHeavyOption = parser.flag("is-heavy")
        .help("defines if this worker handles heavy tasks");
    auto taskId = parser.size_t("task-id")
        .help("services 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 (workerConfig.defined()) {
        configDocPtr.reset(new common::ExtendedXmlDoc(workerConfig));
    } else {
        configDocPtr = common::loadDefaultConfig();
    }

    bool isHeavy = isHeavyOption.defined() && isHeavyOption;
    auto validatorConfig = [&] {
        return validator::ValidatorConfig(
            configDocPtr->get<std::string>(validation::EDITOR_CONFIG_XPATH));
    };

    if (taskId.defined()) {
        if (yt) {
            validation::launchInYtWithoutSplit(*configDocPtr, taskId, isHeavy);
        } else {
            validation::launchLocalWithoutSplit(*configDocPtr, validatorConfig(), taskId, isHeavy);
        }
        return EXIT_SUCCESS;
    }

    worker::Options workerOpts;
    if (grinderConfig.defined()) {
        workerOpts.setConfigLocation(grinderConfig);
    }

    auto handler = [&](const worker::Task& task) {
        if (yt) {
            validation::launchInYtWithSplit(*configDocPtr, validatorConfig(), task, isHeavy);
        } else {
            validation::launchLocalWithSplit(*configDocPtr, validatorConfig(), task, isHeavy);
        }
    };

    workerOpts.on(taskType(isHeavy), handler)
        .setConcurrencyLevel(isHeavy && !yt ? HEAVY_LOCAL_WORKER_CONCURRENCY : WORKER_CONCURRENCY);

    worker::run(workerOpts);

    return EXIT_SUCCESS;
}
catch (const maps::Exception& e) {
    ERROR() << "Worker failed: " << e;
    return EXIT_FAILURE;
}
catch (const std::exception& e) {
    ERROR() << "Worker failed: " << e.what();
    return EXIT_FAILURE;
}
