#include <maps/wikimap/mapspro/services/mrc/eye/lib/recognition_task/include/manager.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/recognition_task/impl/metadata.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/recognition_task/impl/utils.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/recognition_task/impl/traffic_light_handler.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/recognition_task/impl/house_number_handler.h>

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

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

#include <maps/libs/log8/include/log8.h>

namespace maps::mrc::eye {

RecognitionTaskManager::RecognitionTaskManager(const RecognitionTaskManagerConfig& config)
    : BaseMrcWorkerWithConfig(config)
    , typeToHandler_{
          {
              db::eye::RecognitionType::DetectTrafficLight,
              handleTrafficLightRecognitions
          },
          {
              db::eye::RecognitionType::DetectHouseNumber,
              handleHouseNumberRecognitions
          },
      }
{
    REQUIRE(isValid(config_), "Invalid config");
}

size_t RecognitionTaskManager::process(
    pqxx::transaction_base& txn,
    db::eye::Recognitions batch)
{
    size_t updatesNumber = 0;

    for (auto&& [taskType, recognitions] : splitRecognitionsByType(std::move(batch))) {
        updatesNumber += typeToHandler_.at(taskType)(txn, *(config_.frameUrlResolver), std::move(recognitions));
    }

    return updatesNumber;
}

void RecognitionTaskManager::processBatch(const db::TIds& recognitionIds) {
    auto lock = lockIfNeed();

    auto readTxn = getSlaveTxn();
    db::eye::Recognitions recognitions = getUnprocessedRecognitions(
        *readTxn,
        config_.recognitionTypes,
        recognitionIds
    );

    auto writeTxn = getMasterWriteTxn(*(config_.mrc.pool));
    process(*writeTxn, std::move(recognitions));

    commitIfNeed(*writeTxn);
}

bool RecognitionTaskManager::processBatchInLoopMode(size_t batchSize) {
    auto lock = lockIfNeed();

    Batch batch;
    { // read input batch
        auto readTxn = getSlaveTxn();
        auto metadata = recognitionTaskManagerMetadata(*readTxn);

        batch = getNewRecognitionsBatch(
            *readTxn,
            config_.recognitionTypes,
            metadata.getTxnId(),
            batchSize
        );

        INFO() << "Batch size " << batch.recognitions.size()
               << " [" << batch.beginTxnId << ", " << batch.endTxnId << "]";
    }

    auto writeTxn = getMasterWriteTxn(*(config_.mrc.pool));

    process(*writeTxn, std::move(batch.recognitions));

    auto metadata = recognitionTaskManagerMetadata(*writeTxn);
    metadata.updateTxnId(batch.endTxnId);
    metadata.updateTime();

    commitIfNeed(*writeTxn);

    return batch.beginTxnId != batch.endTxnId;
}

} // namespace maps::mrc::eye
