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

#include "detection_relations.h"
#include "eye/recognition.h"

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

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

#include <map>

namespace maps::mrc::eye {

namespace {

db::eye::RecognitionTypes toRecognitionType(const db::eye::DetectionTypes& types)
{
    db::eye::RecognitionTypes result;

    for (auto type: types) {
        result.push_back(db::eye::toRecognitionType(type));
    }

    return result;
}

std::pair<db::TId, db::TId> getInterval(
        pqxx::transaction_base& txn,
        db::TId beginTxnId,
        const db::eye::RecognitionTypes& types,
        size_t limit)
{
    db::TId endTxnId = beginTxnId;

    namespace table = db::eye::table;

    const auto recognitionTxnIds = db::eye::RecognitionGateway(txn).loadTxnIds(
        table::Recognition::txnId >= beginTxnId
            and table::Recognition::type.in(types)
            and table::Recognition::value.isNotNull()
            and table::Recognition::frameId == table::Frame::id
            and table::Recognition::orientation == table::Frame::orientation
            and not table::Frame::deleted,
        sql_chemistry::limit(limit).orderBy(table::Recognition::txnId)
    );

    const auto deletedFrameTxnIds = db::eye::FrameGateway(txn).loadTxnIds(
        table::Frame::txnId >= beginTxnId
            and table::Frame::deleted,
        sql_chemistry::limit(limit).orderBy(table::Frame::txnId)
    );

    db::TIds txnIds;
    txnIds.reserve(recognitionTxnIds.size() + deletedFrameTxnIds.size());
    txnIds.insert(txnIds.end(), recognitionTxnIds.begin(), recognitionTxnIds.end());
    txnIds.insert(txnIds.end(), deletedFrameTxnIds.begin(), deletedFrameTxnIds.end());
    std::sort(txnIds.begin(), txnIds.end());

    if (txnIds.size() > limit) {
        endTxnId = txnIds[limit - 1] + 1;
    } else if (!txnIds.empty()) {
        endTxnId = txnIds.back() + 1;
    }

    return {beginTxnId, endTxnId};
}

db::IdTo<db::eye::DetectionRelations> loadRelationsByGroupId(
    pqxx::transaction_base& txn,
    const db::IdTo<db::TId>& detectionIdToGroupId)
{
    db::TIds detectionIds;
    for (const auto& [detectionId, _] : detectionIdToGroupId) {
        detectionIds.push_back(detectionId);
    }

    db::eye::DetectionRelations relations = db::eye::DetectionRelationGateway(txn).load(
        db::eye::table::DetectionRelation::masterDetectionId.in(detectionIds) ||
        db::eye::table::DetectionRelation::slaveDetectionId.in(detectionIds)
    );
    db::IdTo<db::eye::DetectionRelations> relationsByGroupId;
    for (const auto& relation : relations) {
        db::TId groupId = detectionIdToGroupId.at(relation.masterDetectionId());
        relationsByGroupId[groupId].push_back(relation);
    }

    return relationsByGroupId;
}

// Возвращает два отображения
//   first -  groupId -> множество детекций
//   second - detectionId -> groupId
std::pair<db::IdTo<db::eye::Detections>, db::IdTo<db::TId>>
loadDetectionToGroupMaps(
    pqxx::transaction_base& txn,
    const std::map<std::pair<db::TId, db::eye::DetectionType>, db::eye::DetectionGroup>& frameIdTypeToGroup)
{
    db::TIds groupIds;
    for (const auto& [_, group] : frameIdTypeToGroup) {
        groupIds.push_back(group.id());
    }

    db::eye::Detections detections = db::eye::DetectionGateway(txn).load(
        db::eye::table::Detection::groupId.in(groupIds)
    );
    db::IdTo<db::eye::Detections> detectionsByGroupId;
    db::IdTo<db::TId> detectionIdToGroupId;
    for (const auto& detection : detections) {
        detectionsByGroupId[detection.groupId()].push_back(detection);
        detectionIdToGroupId[detection.id()] = detection.groupId();
    }

    return {
        std::move(detectionsByGroupId),
        std::move(detectionIdToGroupId)
    };
}

std::map<std::pair<db::TId, db::eye::DetectionType>, db::eye::DetectionGroup>
loadFrameIdTypeToGroup(
    pqxx::transaction_base& txn,
    const db::TIds& frameIds,
    const db::eye::DetectionTypes& detectionTypes = {})
{
    sql_chemistry::FiltersCollection filter(sql_chemistry::op::Logical::And);
    filter.add(db::eye::table::DetectionGroup::frameId.in(frameIds));
    if (!detectionTypes.empty()) {
        filter.add(db::eye::table::DetectionGroup::type.in(detectionTypes));
    }

    auto groups = db::eye::DetectionGroupGateway(txn).load(filter);

    std::map<std::pair<db::TId, db::eye::DetectionType>, db::eye::DetectionGroup> frameIdTypeToGroup;
    for (const auto& group : groups) {
        frameIdTypeToGroup.emplace(std::make_pair(group.frameId(), group.type()), group);
    }

    return frameIdTypeToGroup;
}

std::map<std::pair<db::TId, db::eye::RecognitionType>, db::eye::Recognitions>
makeFrameIdTypeToRecognitions(const db::eye::Recognitions& recognitions)
{
    std::map<std::pair<db::TId, db::eye::RecognitionType>, db::eye::Recognitions> frameIdTypeToRecognitions;

    for (const auto& recognition : recognitions) {
        frameIdTypeToRecognitions[{recognition.frameId(), recognition.type()}].push_back(recognition);
    }

    return frameIdTypeToRecognitions;
}

BatchItem makeBatchItem(
    db::TId frameId,
    db::eye::RecognitionType recognitionType,
    const db::eye::Recognitions& recognitions,
    const db::IdTo<db::eye::Frame>& frameById,
    const std::map<std::pair<db::TId, db::eye::DetectionType>, db::eye::DetectionGroup>& frameIdTypeToGroup,
    const db::IdTo<db::eye::Detections>& detectionsByGroupId,
    const db::IdTo<db::eye::DetectionRelations>& relationsByGroupId)
{
    auto detectionType = toDetectionType(recognitionType);
    REQUIRE(detectionType.has_value(), "Incorrect recognition type " << recognitionType);

    std::optional<db::eye::DetectionGroup> group;
    db::eye::Detections detections;
    db::eye::DetectionRelations relations;
    auto groupIt = frameIdTypeToGroup.find(std::make_pair(frameId, detectionType.value()));
    if (groupIt != frameIdTypeToGroup.end()) {
        group = groupIt->second;

        auto detectionsIt = detectionsByGroupId.find(group->id());
        if (detectionsIt != detectionsByGroupId.end()) {
            for (const auto& detection : detectionsIt->second) {
                if (detection.deleted()) {
                    continue;
                }

                detections.push_back(detection);
            }
        }

        auto relationsIt = relationsByGroupId.find(group->id());
        if (relationsIt != relationsByGroupId.end()) {
            relations = relationsIt->second;
        }
    }

    return {
        frameById.at(frameId),
        recognitions,
        group,
        detections,
        relations
    };
}

BatchItems getBatchItemsByRecognitions(
    pqxx::transaction_base& txn,
    const db::eye::Recognitions& recognitions,
    const db::eye::DetectionTypes& detectionTypes)
{
    db::TIds frameIds;
    for (const auto& recognition : recognitions) {
        frameIds.push_back(recognition.frameId());
    }

    auto frameIdTypeToRecognitions = makeFrameIdTypeToRecognitions(recognitions);
    auto frameById = byId(db::eye::FrameGateway(txn).loadByIds(frameIds));
    auto frameIdTypeToGroup = loadFrameIdTypeToGroup(txn, frameIds, detectionTypes);
    auto [detectionsByGroupId, detectionIdToGroupId] = loadDetectionToGroupMaps(txn, frameIdTypeToGroup);
    auto relationsByGroupId = loadRelationsByGroupId(txn, detectionIdToGroupId);

    BatchItems items;
    for (const auto& [frameIdType, recognitions] : frameIdTypeToRecognitions) {
        const auto& [frameId, recognitionType] = frameIdType;

        auto detectionType = toDetectionType(recognitionType);
        if (!detectionType.has_value()) {
            continue;
        }

        items.push_back(
            makeBatchItem(frameId, recognitionType, recognitions,
                frameById,
                frameIdTypeToGroup,
                detectionsByGroupId,
                relationsByGroupId
            )
        );
    }

    return items;
}

BatchItems getBatchItemsByRecognitions(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    db::TId endTxnId,
    const db::eye::DetectionTypes& detectionTypes)
{
    const auto recognitionTypes = toRecognitionType(detectionTypes);

    const auto recognitions = db::eye::RecognitionGateway(txn).load(
        db::eye::table::Recognition::txnId >= beginTxnId && db::eye::table::Recognition::txnId < endTxnId
            && db::eye::table::Recognition::type.in(recognitionTypes)
            && db::eye::table::Recognition::value.isNotNull()
            && db::eye::table::Recognition::frameId == db::eye::table::Frame::id
            && db::eye::table::Recognition::orientation == db::eye::table::Frame::orientation
            && !db::eye::table::Frame::deleted
    );

    return getBatchItemsByRecognitions(txn, recognitions, detectionTypes);
}

BatchItems getBatchItemsByRecognitions(
    pqxx::transaction_base& txn,
    const db::TIds& frameIds,
    const db::eye::DetectionTypes& detectionTypes)
{
    const auto recognitionTypes = toRecognitionType(detectionTypes);

    const auto recognitions = db::eye::RecognitionGateway(txn).load(
        db::eye::table::Recognition::frameId.in(frameIds)
            && db::eye::table::Recognition::type.in(recognitionTypes)
            && db::eye::table::Recognition::value.isNotNull()
            && db::eye::table::Recognition::frameId == db::eye::table::Frame::id
            && db::eye::table::Recognition::orientation == db::eye::table::Frame::orientation
            && !db::eye::table::Frame::deleted
    );

    return getBatchItemsByRecognitions(txn, recognitions, detectionTypes);
}


BatchItems getBatchItemsByFrames(
    pqxx::transaction_base& txn,
    const db::IdTo<db::eye::Frame>& deletedFrameById)
{
    db::TIds frameIds;
    for (const auto& [frameId, _] : deletedFrameById) {
        frameIds.push_back(frameId);
    }

    const auto recognitions = db::eye::RecognitionGateway(txn).load(
        db::eye::table::Recognition::frameId.in(frameIds)
            && db::eye::table::Recognition::value.isNotNull()
            && db::eye::table::Recognition::frameId == db::eye::table::Frame::id
            && db::eye::table::Recognition::orientation == db::eye::table::Frame::orientation
    );

    auto frameIdTypeToRecognitions = makeFrameIdTypeToRecognitions(recognitions);
    auto frameIdTypeToGroup = loadFrameIdTypeToGroup(txn, frameIds);
    auto [detectionsByGroupId, detectionIdToGroupId] = loadDetectionToGroupMaps(txn, frameIdTypeToGroup);
    auto relationsByGroupId = loadRelationsByGroupId(txn, detectionIdToGroupId);

    BatchItems items;
    for (const auto& [frameIdType, recognitions] : frameIdTypeToRecognitions) {
        const auto& [frameId, recognitionType] = frameIdType;

        auto detectionType = toDetectionType(recognitionType);
        if (!detectionType.has_value()) {
            continue;
        }

        items.push_back(
            makeBatchItem(frameId, recognitionType, recognitions,
                deletedFrameById,
                frameIdTypeToGroup,
                detectionsByGroupId,
                relationsByGroupId
            )
        );
    }

    return items;
}

BatchItems getBatchItemsByDeletedFrames(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    db::TId endTxnId)
{
    auto deletedFrameById = byId(
        db::eye::FrameGateway(txn).load(
            db::eye::table::Frame::txnId >= beginTxnId and db::eye::table::Frame::txnId < endTxnId
                && db::eye::table::Frame::deleted
        )
    );

    return getBatchItemsByFrames(txn, deletedFrameById);
}

BatchItems getBatchItemsByDeletedFrames(
    pqxx::transaction_base& txn,
    const db::TIds& frameIds)
{
    auto deletedFrameById = byId(
        db::eye::FrameGateway(txn).load(
            db::eye::table::Frame::id.in(frameIds)
                && db::eye::table::Frame::deleted
        )
    );

    return getBatchItemsByFrames(txn, deletedFrameById);
}

} // namespace

Batch getNewBatch(
    pqxx::transaction_base& txn,
    db::TId txnId,
    const db::eye::DetectionTypes& detectionTypes,
    size_t limit)
{
    const auto recognitionTypes = toRecognitionType(detectionTypes);
    const auto [beginTxnId, endTxnId] = getInterval(txn, txnId, recognitionTypes, limit);

    auto items = getBatchItemsByRecognitions(txn, beginTxnId, endTxnId, detectionTypes);
    auto itemsByDeletedFrames = getBatchItemsByDeletedFrames(txn, beginTxnId, endTxnId);

    items.insert(items.end(),
        std::make_move_iterator(itemsByDeletedFrames.begin()),
        std::make_move_iterator(itemsByDeletedFrames.end())
    );

    return {
        std::move(items),
        beginTxnId,
        endTxnId
    };
}

BatchItems getBatchItems(
    pqxx::transaction_base& txn,
    const db::TIds& frameIds,
    const db::eye::DetectionTypes& detectionTypes)
{
    auto items = getBatchItemsByRecognitions(txn, frameIds, detectionTypes);
    auto itemsByDeletedFrames = getBatchItemsByDeletedFrames(txn, frameIds);

    items.insert(items.end(),
        std::make_move_iterator(itemsByDeletedFrames.begin()),
        std::make_move_iterator(itemsByDeletedFrames.end())
    );

    return items;
}

} // namespace maps::mrc::eye
