#include <maps/wikimap/mapspro/services/mrc/eye/lib/detect_road_marking/impl/batch.h>

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

namespace maps::mrc::eye {

namespace {

// Формирует отображение frameId в recognition панели машины
//
// При этом на одном frame'е может быть больше одного recognition'а
// панели, поэтому выбирает только один с наибольшим приоритетом
// Приоритет основывается на source и version у recognition'а
db::IdTo<db::eye::Recognition> makePanelByFrameId(
    db::eye::Recognitions&& recognitions)
{
    db::IdTo<db::eye::Recognition> panelByFrameId;
    for (db::eye::Recognition& recognition : recognitions) {
        const db::TId frameId = recognition.frameId();

        auto it = panelByFrameId.find(frameId);
        if (panelByFrameId.end() != it) {
            if (it->second.orderKey() < recognition.orderKey()) {
                it->second = std::move(recognition);
            }
        } else {
            panelByFrameId.emplace(frameId, std::move(recognition));
        }
    }

    return panelByFrameId;
}

// Загружает из базы recognition'ы панелей машин на той же ориентации,
// что и соответствующий frame и формирует из них отображение из frameId
// в recognition с наибольшим приоритетом
db::IdTo<db::eye::Recognition> loadPanelByFrameId(
    pqxx::transaction_base& txn,
    const db::IdTo<db::eye::Frame>& frameById)
{
    db::TIds frameIds;
    for (const auto& [frameId, _] : frameById) {
        frameIds.push_back(frameId);
    }

    db::eye::Recognitions recognitions = db::eye::RecognitionGateway(txn).load(
        db::eye::table::Recognition::type == db::eye::RecognitionType::DetectPanel &&
        db::eye::table::Recognition::frameId.in(frameIds) &&
        db::eye::table::Recognition::frameId == db::eye::table::Frame::id &&
        db::eye::table::Recognition::orientation == db::eye::table::Frame::orientation
    );

    return makePanelByFrameId(std::move(recognitions));
}

} // namespace

DetectRoadMarkingInputData loadDetectRoadMarkingInputData(
    pqxx::transaction_base& txn,
    const db::TIds& frameIds)
{
    db::IdTo<db::eye::Frame> frameById = byId(
        db::eye::FrameGateway(txn).load(
            db::eye::table::Frame::id.in(frameIds) &&
            !db::eye::table::Frame::deleted
        )
    );
    db::IdTo<db::eye::Recognition> panelByFrameId = loadPanelByFrameId(txn, frameById);

    return {
        std::move(frameById),
        std::move(panelByFrameId)
    };
}

namespace {

// Находит txn_id конца нового отрезка в очереди frame'ов
//  * beginTxnId - начало очереди frame'ов
//  * limit - количество элементов загружаемое из базы,
//    при этом из базы может быть прочитано больше элементов,
//    если все они имеют одинаковый txn_id
db::TId getEndFrameTxnId(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    size_t limit)
{
    db::TId endTxnId = beginTxnId;

    db::TIds frameTxnIds = db::eye::FrameGateway(txn).loadTxnIds(
        db::eye::table::Frame::txnId >= beginTxnId,
        sql_chemistry::limit(limit).orderBy(db::eye::table::Frame::txnId)
    );
    if (!frameTxnIds.empty()) {
        endTxnId = std::max(endTxnId, frameTxnIds.back() + 1);
    }

    return endTxnId;
}

// Находит txn_id конца нового отрезка в очереди recognition'ов
//  * beginTxnId - начало очереди recognition'ов
//  * limit - количество элементов загружаемое из базы,
//    при этом из базы может быть прочитано больше элементов,
//    если все они имеют одинаковый txn_id
db::TId getEndPanelTxnId(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    size_t limit)
{
    db::TId endTxnId = beginTxnId;

    db::TIds recognitionTxnIds = db::eye::RecognitionGateway(txn).loadTxnIds(
        db::eye::table::Recognition::txnId >= beginTxnId,
        sql_chemistry::limit(limit).orderBy(db::eye::table::Recognition::txnId)
    );
    if (!recognitionTxnIds.empty()) {
        endTxnId = std::max(endTxnId, recognitionTxnIds.back() + 1);
    }

    return endTxnId;
}

// Загружает набор входных данных из очереди по frame'ам
DetectRoadMarkingInputData loadDataByFrames(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    db::TId endTxnId)
{
    db::IdTo<db::eye::Frame> frameById = byId(
        db::eye::FrameGateway(txn).load(
            db::eye::table::Frame::txnId >= beginTxnId &&
            db::eye::table::Frame::txnId < endTxnId &&
            !db::eye::table::Frame::deleted
        )
    );
    db::IdTo<db::eye::Recognition> panelByFrameId = loadPanelByFrameId(txn, frameById);

    return {
        std::move(frameById),
        std::move(panelByFrameId)
    };
}

// Загружает набор входных данных из очереди по recognition'ам
DetectRoadMarkingInputData loadDataByPanels(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    db::TId endTxnId)
{
    db::eye::Recognitions recognitions = db::eye::RecognitionGateway(txn).load(
        db::eye::table::Recognition::type == db::eye::RecognitionType::DetectPanel &&
        db::eye::table::Recognition::txnId >= beginTxnId &&
        db::eye::table::Recognition::txnId < endTxnId
    );

    db::TIds frameIds;
    for (const auto& recognition : recognitions) {
        frameIds.push_back(recognition.frameId());
    }

    return loadDetectRoadMarkingInputData(txn, frameIds);
}

} // namespace


DetectRoadMarkingBatch loadDetectRoadMarkingBatch(
    pqxx::transaction_base& txn,
    db::TId beginFrameTxnId,
    db::TId beginPanelTxnId,
    size_t limit)
{
    db::TId endFrameTxnId = getEndFrameTxnId(txn, beginFrameTxnId, limit);
    db::TId endPanelTxnId = getEndPanelTxnId(txn, beginPanelTxnId, limit);

    auto data = loadDataByFrames(txn, beginFrameTxnId, endFrameTxnId);
    auto panelData = loadDataByPanels(txn, beginPanelTxnId, endPanelTxnId);

    data.frameById.merge(panelData.frameById);
    data.panelByFrameId.merge(panelData.panelByFrameId);

    return {
        beginFrameTxnId, endFrameTxnId,
        beginPanelTxnId, endPanelTxnId,
        std::move(data)
    };
}

} // namespace maps::mrc::eye
