#include <maps/wikimap/mapspro/services/mrc/eye/lib/detect_panel/include/detect_panel.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/detect_panel/include/metadata.h>
#include <maps/wikimap/mapspro/services/mrc/eye/lib/detect_panel/include/sequence.h>

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

#include <maps/wikimap/mapspro/services/mrc/libs/yt/include/io.h>
#include <maps/wikimap/mapspro/services/mrc/libs/yt/include/schema.h>

#include <maps/wikimap/mapspro/services/mrc/libs/common/include/pg_locks.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/sql_chemistry/include/order.h>

#include <algorithm>
#include <utility>

namespace maps::mrc::eye {

namespace {

db::TIdSet toSet(const db::TIds& frameIds) { return {frameIds.begin(), frameIds.end()}; }

struct Batch {
    db::TIdSet frameIds;
    db::TId beginTxnId;
    db::TId endTxnId;
};

db::TIdSet keepWithoutRecognition(pqxx::transaction_base& txn, db::TIdSet frameIds)
{
    namespace table = db::eye::table;

    if (frameIds.empty()) {
        return frameIds;
    }

    const db::TIds recognizedFrameIds = db::eye::FrameGateway(txn).loadIds(
        table::Frame::id.in({frameIds.begin(), frameIds.end()})
            and table::Frame::id == table::Recognition::frameId
            and table::Frame::orientation == table::Recognition::orientation
            and table::Recognition::type == db::eye::RecognitionType::DetectPanel
    );

    for (auto frameId: recognizedFrameIds) {
        frameIds.erase(frameId);
    }

    return frameIds;
}

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

    return {toSet(batch.ids), batch.beginTxnId, batch.endTxnId};
}

} // namespace

Partitions DetectPanel::load(pqxx::transaction_base& txn, const db::TIdSet& frameIds)
{
    return makePartitions(
        makeFrameSequences(txn, frameIds, config_.margin, std::chrono::seconds(60)),
        config_.yt.partitionSize,
        config_.margin
    );
}

size_t DetectPanel::save(
        pqxx::transaction_base& txn,
        const db::TIdSet& toSaveIds,
        const DetectedPanels& detections)
{
    db::eye::Recognitions recognitions;

    for (const auto& [frame, restBox]: detections) {
        if (not toSaveIds.count(frame.id())) {
            continue;
        }

        recognitions.emplace_back(
            frame.id(),
            frame.orientation(),
            db::eye::RecognitionType::DetectPanel,
            db::eye::RecognitionSource::Model,
            version,
            db::eye::DetectedPanel{restBox}
        );
    }

    if (not recognitions.empty()) {
        db::eye::RecognitionGateway(txn).upsertx(recognitions);
    }

    return recognitions.size();
}

void DetectPanel::processBatch(const db::TIds& frameIds)
{
    const auto lock = lockIfNeed();

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

    const db::TIdSet idSet = toSet(frameIds);
    INFO() << "Frames " << idSet.size();

    const Partitions partitions = load(*getSlaveTxn(), idSet);
    INFO() << "Frame partitions " <<  partitions.frames.size() << ":" <<  partitions.partitionNumber;

    INFO() << "Run detection...";
    const DetectedPanels detections = detectPanel(config_, partitions);

    auto txn = getMasterWriteTxn(*(config_.mrc.pool));
    const size_t count = save(*txn, idSet, detections);
    INFO() << "Save " <<  count << " detections";

    commitIfNeed(*txn);
}

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

    Batch batch;
    Partitions partitions;

    {   // load and free txn
        auto txn = getSlaveTxn();

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

        batch.frameIds = keepWithoutRecognition(*txn, std::move(batch.frameIds));

        partitions = load(*txn, batch.frameIds);
    }

    INFO() << "Frame partitions " <<  partitions.frames.size() << ":" <<  partitions.partitionNumber;

    INFO() << "Run detection...";
    const DetectedPanels detections = detectPanel(config_, partitions);

    auto txn = getMasterWriteTxn(*(config_.mrc.pool));
    const size_t count = save(*txn, batch.frameIds, detections);
    INFO() << "Save " <<  count << " detections";

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

    commitIfNeed(*txn);

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

} // namespace maps::mrc::eye
