#pragma once

#include <maps/wikimap/mapspro/services/mrc/eye/lib/detect_object/include/config.h>

#include <maps/wikimap/mapspro/services/mrc/eye/lib/common/include/base_worker.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/common/include/metadata.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/common.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/eye/frame_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/eye/recognition_gateway.h>

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

#include <pqxx/pqxx>

#include <functional>
#include <optional>
#include <string>


namespace maps::mrc::eye {

typedef MetadataManager(*MetadataFactory)(pqxx::transaction_base&);

typedef db::eye::Recognitions(*SimpleDetectObjectFunctor)(
        const SimpleDetectObjectConfig&,
        const db::eye::Frames&);

// keep frames that don't have any corresponding recognitions
db::TIds keepWithoutRecognition(
        pqxx::transaction_base& txn,
        db::eye::RecognitionType type,
        int16_t version,
        db::TIds frameIds);

void updateRecognitions(
        pqxx::transaction_base& txn,
        db::eye::RecognitionType type,
        int16_t version,
        db::eye::Recognitions& recognitions);

inline db::eye::Frames loadFrames(pqxx::transaction_base& txn, const db::TIds& frameIds)
{
    return db::eye::FrameGateway(txn).loadByIds(frameIds);
}

template<
    SimpleDetectObjectFunctor detect,
    MetadataFactory getMetadata,
    db::eye::RecognitionType type,
    int16_t version,
    common::LockId lockId>
class SimpleDetectObject: public BaseMrcWorkerWithConfig<SimpleDetectObjectConfig, lockId> {

    using Base = BaseMrcWorkerWithConfig<SimpleDetectObjectConfig, lockId>;

public:
    using Base::Base;

    static constexpr int16_t VERSION = version;

    void processBatch(const db::TIds& frameIds) override
    {
        const auto lock = this->lockIfNeed();

        if (frameIds.empty()) {
            INFO() << "Stop, no frames!";
            return;
        }

        db::eye::Frames frames;
        {   // load and free txn
            auto txn = getSlaveTxn(*(this->config_.mrc.pool), this->token());

            if (this->config_.rework) {
                frames = loadFrames(*txn, frameIds);
            } else {
                frames = loadFrames(*txn, keepWithoutRecognition(*txn, type, version, frameIds));
            }
        }
        INFO() << "Load " << frames.size() << " frames";

        db::eye::Recognitions recognitions = detect(this->config_, frames);

        auto txn = getMasterWriteTxn(*(this->config_.mrc.pool));
        if (this->config_.rework) {
            updateRecognitions(*txn, type, version, recognitions);
        } else {
            db::eye::RecognitionGateway(*txn).insertx(recognitions);
        }
        INFO() << "Save " <<  recognitions.size() << " recognitions";

        this->commitIfNeed(*txn);
    }


    bool processBatchInLoopMode(size_t batchSize) override
    {
        const auto lock = this->lockIfNeed();

        Batch batch;
        db::eye::Frames frames;

        {   // load and free txn
            auto txn = getSlaveTxn(*(this->config_.mrc.pool), this->token());

            batch = getNewBatch(*txn, batchSize);
            INFO() << "Batch size " << batch.ids.size()
                << " [" << batch.beginTxnId << ", " << batch.endTxnId << ")";

            if (this->config_.rework) {
                frames = loadFrames(*txn, batch.ids);
            } else {
                frames = loadFrames(*txn, keepWithoutRecognition(*txn, type, version, batch.ids));
            }
        }

        db::eye::Recognitions recognitions = detect(this->config_, frames);

        auto txn = getMasterWriteTxn(*(this->config_.mrc.pool));
        if (this->config_.rework) {
            updateRecognitions(*txn, type, version, recognitions);
        } else {
            db::eye::RecognitionGateway(*txn).insertx(recognitions);
        }
        INFO() << "Save " <<  recognitions.size() << " detections";

        auto metadata = getMetadata(*txn);
        metadata.updateTxnId(batch.endTxnId);
        metadata.updateTime();

        this->commitIfNeed(*txn);

        return batch.ids.size() > 0;
    }

private:
    using Batch = db::eye::FrameGateway::Batch;

    Batch getNewBatch(pqxx::transaction_base& txn, size_t size)
    {
        return db::eye::FrameGateway(txn).loadBatch(
            getMetadata(txn).getTxnId(),
            size,
            not db::eye::table::Frame::deleted
        );
    }
};

} // namespace maps::mrc::eye
