#include "batch.h"
#include "db.h"

#include <maps/wikimap/mapspro/services/mrc/eye/lib/common/include/util.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/wikimap/mapspro/services/mrc/libs/db/include/eye/verified_detection_pair_match_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/eye/verified_detection_missing_on_frame_gateway.h>

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

#include <maps/libs/sql_chemistry/include/exists.h>

#include <algorithm>

namespace maps::mrc::eye {

namespace {

db::TIds getDetectionGroupTxnIds(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    size_t limit)
{
    return db::eye::DetectionGroupGateway{txn}.loadTxnIds(
        db::eye::table::DetectionGroup::txnId >= beginTxnId,
        sql_chemistry::limit(limit)
            .orderBy(db::eye::table::DetectionGroup::txnId)
    );
}

db::TIds getLocationTxnIds(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    size_t limit)
{
    return db::eye::FrameLocationGateway{txn}
        .loadTxnIds(
            db::eye::table::FrameLocation::txnId >= beginTxnId,
            sql_chemistry::limit(limit)
                .orderBy(db::eye::table::FrameLocation::txnId)
    );
}

db::TIds getVerifiedDetectionPairMatchTxnIds(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    size_t limit)
{
    return db::eye::VerifiedDetectionPairMatchGateway{txn}
        .loadTxnIds(
            db::eye::table::VerifiedDetectionPairMatch::txnId >= beginTxnId
            && db::eye::table::VerifiedDetectionPairMatch::approved.isNotNull(),
            sql_chemistry::limit(limit)
                .orderBy(db::eye::table::VerifiedDetectionPairMatch::txnId)
    );
}

db::TIds getVerifiedDetectionMissingOnFrameTxnIds(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    size_t limit)
{
    return db::eye::VerifiedDetectionMissingOnFrameGateway{txn}
        .loadTxnIds(
            db::eye::table::VerifiedDetectionMissingOnFrame::txnId >= beginTxnId
            && db::eye::table::VerifiedDetectionMissingOnFrame::isVisible ==
                db::eye::VerifiedDetectionMissingOnFrameIsVisible::No
            && db::eye::table::VerifiedDetectionMissingOnFrame::missingReason ==
                db::eye::VerifiedDetectionMissingOnFrameMissingReason::Missing,
            sql_chemistry::limit(limit)
                .orderBy(db::eye::table::VerifiedDetectionMissingOnFrame::txnId)
    );
}

db::TId getEndTxnId(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    size_t limit)
{
    if (limit == 0) {
        return beginTxnId;
    }

    auto detectionGroupTxnIds = getDetectionGroupTxnIds(txn, beginTxnId, limit);
    auto locationTxnIds = getLocationTxnIds(txn, beginTxnId, limit);
    auto vdpm = getVerifiedDetectionPairMatchTxnIds(txn, beginTxnId, limit);
    auto vdmof = getVerifiedDetectionMissingOnFrameTxnIds(txn, beginTxnId, limit);

    db::TIds txnIds;
    txnIds.reserve(detectionGroupTxnIds.size() + locationTxnIds.size() + vdpm.size());
    txnIds.insert(txnIds.end(), detectionGroupTxnIds.begin(), detectionGroupTxnIds.end());
    txnIds.insert(txnIds.end(), locationTxnIds.begin(), locationTxnIds.end());
    txnIds.insert(txnIds.end(), vdpm.begin(), vdpm.end());
    txnIds.insert(txnIds.end(), vdmof.begin(), vdmof.end());
    std::sort(txnIds.begin(), txnIds.end());

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

db::IdTo<db::TId> mergePreservingMaxValue(
    db::IdTo<db::TId>&& lhs,
    db::IdTo<db::TId>&& rhs)
{
    db::IdTo<db::TId> result;
    for (const auto [id, txnId] : lhs) {
        if (rhs.count(id)) {
            result.emplace(id, std::max(rhs.at(id), txnId));
        } else {
            result.emplace(id, txnId);
        }
    }

    for (const auto [id, txnId] : rhs) {
        if (!result.count(id)) {
            result.emplace(id, txnId);
        }
    }
    return result;
}

db::IdTo<db::TId> loadUpdatedDetectionIdToTxnId(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    db::TId endTxnId,
    const db::eye::DetectionTypes& detectionTypes)
{
    namespace table = db::eye::table;

    const auto values =  db::eye::DetectionGateway(txn).loadJoined<db::TId>(
        std::make_tuple(table::DetectionGroup::txnId),
        table::DetectionGroup::txnId >= beginTxnId
            && table::DetectionGroup::txnId < endTxnId
            && table::Detection::groupId == table::DetectionGroup::id
            && table::DetectionGroup::type.in(detectionTypes)
    );

    db::IdTo<db::TId> detectionIdToTxn;
    for (const auto& [txnId, detection] : values) {
        detectionIdToTxn[detection.id()] = txnId;
    }
    return detectionIdToTxn;
}

db::IdTo<db::TId> loadMovedDetectionIdToTxnId(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    db::TId endTxnId,
    const db::eye::DetectionTypes& detectionTypes)
{
    namespace table = db::eye::table;

    const auto values =  db::eye::DetectionGateway(txn).loadJoined<db::TId>(
        std::make_tuple(table::FrameLocation::txnId),
        table::FrameLocation::txnId >= beginTxnId and table::FrameLocation::txnId < endTxnId
            and table::FrameLocation::frameId == table::Frame::id
            and not table::Frame::deleted
            and table::Frame::id == table::DetectionGroup::frameId
            and table::DetectionGroup::type.in(detectionTypes)
            and table::DetectionGroup::id == table::Detection::groupId
            and table::Detection::txnId < beginTxnId
            and not table::Detection::deleted
    );

    db::IdTo<db::TId> detectionIdToTxn;
    for (const auto& [txnId, detection] : values) {
        detectionIdToTxn[detection.id()] = txnId;
    }
    return detectionIdToTxn;
}

db::IdTo<db::TId> loadDetectionsByVerifiedMatchesIdToTxnId(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    db::TId endTxnId,
    const db::eye::DetectionTypes& detectionTypes)
{
    namespace table = db::eye::table;

    const auto values = db::eye::DetectionGateway(txn).loadJoined<db::TId>(
        std::make_tuple(table::VerifiedDetectionPairMatch::txnId),
        table::VerifiedDetectionPairMatch::txnId >= beginTxnId
            and table::VerifiedDetectionPairMatch::txnId < endTxnId
            and table::VerifiedDetectionPairMatch::approved.isNotNull()
            and (table::VerifiedDetectionPairMatch::detectionId1 == table::Detection::id
                 or table::VerifiedDetectionPairMatch::detectionId2 == table::Detection::id)
            and table::DetectionGroup::id == table::Detection::groupId
            and table::DetectionGroup::type.in(detectionTypes)
            and table::Detection::txnId < beginTxnId
            and not table::Detection::deleted
    );

    db::IdTo<db::TId> detectionGroupIdToTxn;
    for (const auto& [txnId, group] : values) {
        detectionGroupIdToTxn[group.id()] = txnId;
    }
    return detectionGroupIdToTxn;
}

db::IdTo<db::TId>
loadDetectionsByVerifiedMissingOnFrameIdToTxnId(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    db::TId endTxnId,
    const db::eye::DetectionTypes& detectionTypes)
{
    namespace table = db::eye::table;

    const auto values = db::eye::DetectionGateway(txn).loadJoined<db::TId>(
        std::make_tuple(table::VerifiedDetectionMissingOnFrame::txnId),
        table::VerifiedDetectionMissingOnFrame::txnId >= beginTxnId
            and table::VerifiedDetectionMissingOnFrame::txnId < endTxnId
            and table::VerifiedDetectionMissingOnFrame::isVisible ==
                db::eye::VerifiedDetectionMissingOnFrameIsVisible::No
            and table::VerifiedDetectionMissingOnFrame::missingReason ==
                db::eye::VerifiedDetectionMissingOnFrameMissingReason::Missing
            and table::VerifiedDetectionMissingOnFrame::detectionId
                == table::Detection::id
            and table::DetectionGroup::id == table::Detection::groupId
            and table::DetectionGroup::type.in(detectionTypes)
            and table::Detection::txnId < beginTxnId
            and not table::Detection::deleted
    );

    db::IdTo<db::TId> detectionIdToTxn;
    for (const auto& [txnId, detection] : values) {
        detectionIdToTxn[detection.id()] = txnId;
    }
    return detectionIdToTxn;
}


db::TIds loadDetectionGroupIdsForDetections(
    pqxx::transaction_base& txn,
    const db::TIds& detectionIds
) {
    return db::eye::DetectionGroupGateway(txn).loadIds(
        sql_chemistry::existsIn<db::eye::table::Detection>(
            db::eye::table::Detection::id.in(detectionIds)
            && db::eye::table::Detection::groupId ==
                db::eye::table::DetectionGroup::id
        )
    );
}

db::TIds loadEmptyDetectionGroupIds(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    db::TId endTxnId,
    const db::eye::DetectionTypes& detectionTypes
)
{
    return db::eye::DetectionGroupGateway(txn).loadIds(
        db::eye::table::DetectionGroup::type.in(detectionTypes)
            && db::eye::table::DetectionGroup::txnId >= beginTxnId
            && db::eye::table::DetectionGroup::txnId < endTxnId
            && !sql_chemistry::existsIn<db::eye::table::Detection>(
                db::eye::table::Detection::groupId ==
                    db::eye::table::DetectionGroup::id
            ));
}

std::vector<TxnIdDetectionGroupId> loadDetectionGroupIds(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    db::TId endTxnId,
    const db::eye::DetectionTypes& detectionTypes)
{
    auto detectionIdToTxn =  mergePreservingMaxValue(
        loadUpdatedDetectionIdToTxnId(txn, beginTxnId, endTxnId, detectionTypes),
        loadMovedDetectionIdToTxnId(txn, beginTxnId, endTxnId, detectionTypes)
    );

    detectionIdToTxn = mergePreservingMaxValue(
        std::move(detectionIdToTxn),
        loadDetectionsByVerifiedMatchesIdToTxnId(txn, beginTxnId, endTxnId, detectionTypes)
    );

    detectionIdToTxn = mergePreservingMaxValue(
        std::move(detectionIdToTxn),
        loadDetectionsByVerifiedMissingOnFrameIdToTxnId(txn, beginTxnId, endTxnId, detectionTypes)
    );

    const db::TIds detectionGroupIds = loadDetectionGroupIdsForDetections(txn, collectKeys(detectionIdToTxn));
    const db::IdTo<db::TIds> groupIdToDetectionIds = loadGroupIdToDetectionIdsMap(txn, detectionGroupIds);
    const db::TIds emptyDetectionGroupIds = loadEmptyDetectionGroupIds(txn, beginTxnId, endTxnId, detectionTypes);
    std::vector<TxnIdDetectionGroupId> result;
    result.reserve(emptyDetectionGroupIds.size() + groupIdToDetectionIds.size());
    for (auto groupId : emptyDetectionGroupIds) {
        result.push_back({.detectionGroupId = groupId});
    }
    for (const auto& [groupId, detectionIds] : groupIdToDetectionIds) {
        std::vector<TxnIdDetectionId> txnDetectionIds;
        txnDetectionIds.reserve(detectionIds.size());
        for (auto detectionId : detectionIds) {
            auto txnIt = detectionIdToTxn.find(detectionId);
            txnDetectionIds.push_back({
                .txnId  = txnIt != detectionIdToTxn.end() ? txnIt->second : 0,
                .detectionId = detectionId});
        }
        result.push_back({.detectionGroupId = groupId, .txnDetectionIds = std::move(txnDetectionIds)});
    }
    return result;
}

} // namespace

Batch getNewBatch(
    pqxx::transaction_base& txn,
    db::TId beginTxnId,
    size_t limit,
    const db::eye::DetectionTypes& detectionTypes)
{
    const db::TId endTxnId = getEndTxnId(txn, beginTxnId, limit);

    return {
        beginTxnId, endTxnId,
        loadDetectionGroupIds(txn, beginTxnId, endTxnId, detectionTypes)
    };
}

db::TIds extractDetectionGroupIds(const std::vector<TxnIdDetectionGroupId>& txnDetectionGroupIds)
{
    db::TIds result;
    result.reserve(txnDetectionGroupIds.size());
    for (const auto& txnDetectionGroupId : txnDetectionGroupIds) {
        result.push_back(txnDetectionGroupId.detectionGroupId);
    }
    return result;
}

db::TIds extractDetectionIds(const std::vector<TxnIdDetectionId>& txnDetectionIds)
{
    db::TIds result;
    result.reserve(txnDetectionIds.size());
    for (const auto& txnDetectionId : txnDetectionIds) {
        result.push_back(txnDetectionId.detectionId);
    }
    return result;
}

db::TIds extractDetectionIds(const std::vector<struct TxnIdDetectionGroupId>& txnDetectionGroupIds)
{
    db::TIds result;
    for (const auto& [_, txnDetectionIds] : txnDetectionGroupIds) {
        const auto detectionIds = extractDetectionIds(txnDetectionIds);
        result.insert(result.end(), detectionIds.begin(), detectionIds.end());
    }
    return result;
}

} // namespace maps::mrc::eye
