#include <maps/wikimap/mapspro/services/mrc/eye/lib/import_detection/include/importer.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/import_detection/include/metadata.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/import_detection/include/handler.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/wikimap/mapspro/services/mrc/libs/db/include/feature_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/metadata_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/eye/frame_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/eye/feature_gateway.h>

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

namespace maps::mrc::eye {

namespace {

using ImportDetectionHandler
    = std::function<size_t(pqxx::transaction_base&, const db::eye::Frames&)>;

std::map<db::eye::RecognitionType, ImportDetectionHandler> TYPE_TO_HANDLER{
    {
        db::eye::RecognitionType::DetectTrafficLight,
        handleImportTrafficLightDetection
    },
    {
        db::eye::RecognitionType::DetectSign,
        handleImportSignDetection
    },
    {
        db::eye::RecognitionType::DetectHouseNumber,
        handleImportHouseNumberDetection
    },
    {
        db::eye::RecognitionType::DetectRoadMarking,
        handleImportRoadMarkingDetection
    },
};

struct Batch {
    db::eye::Frames frames;
    db::TId beginTxnId;
    db::TId endTxnId;
};

Batch getNewBatch(pqxx::transaction_base& txn, size_t batchSize, db::TId beginTxnId) {
    const auto batch = db::eye::FrameGateway(txn).loadBatch(
        beginTxnId, batchSize, not db::eye::table::Frame::deleted
    );

    db::eye::Frames frames = db::eye::FrameGateway(txn).loadByIds(batch.ids);

    return {frames, batch.beginTxnId, batch.endTxnId};
}

} // namespace

DetectionImporter::DetectionImporter(const ImportDetectionConfig& config)
    : config_(config)
{
    REQUIRE(isValid(config_), "Invalid config");
}

size_t DetectionImporter::processBatch(const db::TIds& frameIds) {
    auto txn = config_.mrc.pool->masterWriteableTransaction();

    db::eye::Frames frames = db::eye::FrameGateway(*txn).loadByIds(frameIds);

    size_t importedNumber = 0;
    for (db::eye::RecognitionType recognitionType : config_.recognitionTypes) {
        INFO() << "Import " << toString(recognitionType) << " detections";

        auto txn = config_.mrc.pool->masterWriteableTransaction();

        importedNumber += TYPE_TO_HANDLER.at(recognitionType)(*txn, frames);

        commitIfNeed(*txn);
    }

    return importedNumber;
}

std::pair<size_t, bool> DetectionImporter::processBatchInLoopMode(size_t batchSize) {
    size_t importedNumber = 0;
    bool isImportCompleted = true;

    for (db::eye::RecognitionType recognitionType : config_.recognitionTypes) {
        INFO() << "Import " << toString(recognitionType) << " detections";

        auto txn = config_.mrc.pool->masterWriteableTransaction();

        db::TId txnId = getImportDetectionTxnId(*txn, recognitionType);
        Batch batch = getNewBatch(*txn, batchSize, txnId);
        INFO() << "Batch [" << batch.beginTxnId << ", " << batch.endTxnId << ")";

        if (!batch.frames.empty()) {
            isImportCompleted = false;
        } else {
            continue;
        }

        importedNumber += TYPE_TO_HANDLER.at(recognitionType)(*txn, batch.frames);

        updateImportDetectionTxnId(*txn, recognitionType, batch.endTxnId);

        commitIfNeed(*txn);
    }

    return std::make_pair(importedNumber, isImportCompleted);
}

void DetectionImporter::commitIfNeed(pqxx::transaction_base& txn) {
    if (config_.mrc.commit) {
        txn.commit();
        INFO() << "Commited!";
    } else {
        INFO() << "Commit aborted due to dry run mode!";
    }
}

} // namespace maps::mrc::eye
