#pragma once

#include <maps/wikimap/mapspro/services/mrc/libs/common/include/pg_locks.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/common.h>

#include <maps/wikimap/mapspro/services/mrc/eye/lib/common/include/txn.h>

#include <maps/libs/common/include/exception.h>
#include <maps/libs/log8/include/log8.h>
#include <maps/libs/pgpool3utils/include/yandex/maps/pgpool3utils/pg_advisory_mutex.h>

#include <chrono>
#include <thread>

namespace maps::mrc::eye {

template<typename PK>
class BaseWorker {
public:
    using PrimaryKey = PK;
    using PrimaryKeys = std::vector<PrimaryKey>;

    virtual void processBatch(const PrimaryKeys&) = 0;
    virtual bool processBatchInLoopMode(size_t /* batch size */) = 0;
    virtual void runInLoopMode(size_t batchSize, std::chrono::minutes timeout) {
        for (;;) {
            INFO() << "Process new batch...";

            bool hasMore = false;
            try {
                hasMore = processBatchInLoopMode(batchSize);
            } catch (const std::exception& e) {
                WARN() << "Failed to process batch: " << FormatCurrentException();
            }

            if (!hasMore) {
                INFO() << "Sleep for timeout " << timeout.count() << " min";
                std::this_thread::sleep_for(timeout);
            }
        }
    }

    virtual ~BaseWorker() = default;
};

template <typename WorkerConfig, common::LockId lockId, typename PK>
class BaseWorkerWithConfig : public BaseWorker<PK> {
public:
    BaseWorkerWithConfig(const WorkerConfig& config)
        : config_(config)
    {
        REQUIRE(isValid(config_), "Invalid config");
    }

protected:
    pgpool3::Pool* pool() { return config_.mrc.pool; };

    std::optional<pgp3utils::PgAdvisoryXactMutex> lockIfNeed() {
        if (config_.mrc.lockFree || not config_.mrc.commit) {
            return std::nullopt;
        }
        pgp3utils::PgAdvisoryXactMutex mutex(
            *pool(),
            static_cast<int64_t>(lockId)
        );
        INFO() << "Wait for lock...";
        mutex.lock();
        INFO() << "Ready!";
        return mutex;
    }

    pgpool3::TransactionHandle getSlaveTxn() {
        return maps::mrc::eye::getSlaveTxn(*(config_.mrc.pool), token());
    }

    void commitIfNeed(pqxx::transaction_base& txn) {
        if (config_.mrc.commit) {
            token_ = commit(txn);
            INFO() << "Commited!";
        } else {
            INFO() << "Skip commit!";
        }
    }

    std::string token() {
        if (token_.empty()) {
            return getNewToken(*(config_.mrc.pool));
        }
        return token_;
    }

    std::string token_;
    WorkerConfig config_;
};

using BaseMrcWorker = BaseWorker<db::TId>;

template<typename WorkerConfig, common::LockId lockId>
class BaseMrcWorkerWithConfig :
    public BaseWorkerWithConfig<WorkerConfig, lockId, db::TId> {
public:
    using BaseWorkerWithConfig<WorkerConfig, lockId, db::TId>::BaseWorkerWithConfig;

protected:
    using BaseWorkerWithConfig<WorkerConfig, lockId, db::TId>::lockIfNeed;
    using BaseWorkerWithConfig<WorkerConfig, lockId, db::TId>::commitIfNeed;
    using BaseWorkerWithConfig<WorkerConfig, lockId, db::TId>::token;
};

} // namespace maps::mrc::eye
