#include "loader.h"

#include <maps/libs/log8/include/log8.h>
#include <maps/libs/sql_chemistry/include/batch_load.h>
#include <maps/wikimap/mapspro/services/mrc/libs/common/include/algorithm/collection.h>
#include <maps/wikimap/mapspro/services/mrc/libs/common/include/algorithm/for_each_batch.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/feature_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/object_in_photo_gateway.h>
#include <maps/wikimap/mapspro/services/mrc/libs/db/include/walk_object_gateway.h>

namespace maps::mrc::export_gen {

namespace {

db::TIds toFeatureIds(const db::Features& features)
{
    auto result = db::TIds{};
    result.reserve(features.size());
    for (const auto& feature : features) {
        result.push_back(feature.id());
    }
    return result;
}

ObjectsInPhotoMap toObjectsInPhotoMap(db::ObjectsInPhoto&& objects)
{
    auto result = ObjectsInPhotoMap{};
    for (auto&& object : objects) {
        auto featureId = object.featureId();
        result[featureId].push_back(std::move(object));
    }
    return result;
}

void forEachWalkObjects(pgpool3::Pool& pool, size_t batchSize, WalkObjectsFn fn)
{
    auto batch = sql_chemistry::BatchLoad<db::table::WalkObject>{batchSize};
    while (batch.next(*pool.slaveTransaction())) {
        fn(std::move(batch.result()));
    }
}

db::TIds loadNewFeatureIds(sql_chemistry::Transaction& txn, db::TId lastTxnId)
{
    auto featureTxns = db::FeatureTransactionGateway{txn}.load(
        db::table::FeatureTransaction::transactionId > lastTxnId);
    auto result = db::TIds{};
    for (const auto& featureTxn : featureTxns) {
        result.push_back(featureTxn.featureId());
    }
    return result;
}

}  // namespace

void DbLoader::forEachPublishedFeatures(size_t batchSize, FeaturesFn fn) const
{
    auto batch = sql_chemistry::BatchLoad<db::table::Feature>{
        batchSize,
        db::table::Feature::isPublished.is(true) ||
            db::table::Feature::shouldBePublished.is(true)};
    while (batch.next(*pool_.slaveTransaction())) {
        auto featureIds = toFeatureIds(batch.result());
        auto objects = db::ObjectInPhotoGateway{*pool_.slaveTransaction()}.load(
            db::table::ObjectInPhotoTable::featureId.in(std::move(featureIds)));
        auto objectsMap = toObjectsInPhotoMap(std::move(objects));
        fn(std::move(batch.result()), std::move(objectsMap), Origin::Db);
    }
}

void DbLoader::forEachWalkObjects(size_t batchSize, WalkObjectsFn fn) const
{
    export_gen::forEachWalkObjects(pool_, batchSize, fn);
}

FbLoader::FbLoader(pgpool3::Pool& pool,
                   const std::string& mrcDatasetDir,
                   const std::string& mrcFeaturesSecretDatasetDir,
                   db::TId lastTxnId,
                   EMappingMode mappingMode)
    : pool_(pool)
    , defaultFeatures_(mrcDatasetDir, mappingMode)
    , secretFeatures_(mrcFeaturesSecretDatasetDir, mappingMode)
    , lastTxnId_(std::min(lastTxnId, defaultFeatures_.lastTxnId()))
{
    REQUIRE(defaultFeatures_.version() == secretFeatures_.version(),
            "inconsistent versions of default features ("
                << defaultFeatures_.version() << "), secret features ("
                << secretFeatures_.version() << ")");
    REQUIRE(lastTxnId_, "no lastTxnId");
    REQUIRE(
        defaultFeatures_.schemaVersion() == CURRENT_FEATURE_SCHEMA_VERSION,
        "old feature schema version");
    INFO() << "FbLoader{.version=" << defaultFeatures_.version()
           << ", .lastTxnId=" << lastTxnId_ << "}";
}

db::TIdSet FbLoader::loadNewPublishedFeatures(size_t batchSize,
                                              const FeaturesFn& fn) const
{
    auto newFeatureIds =
        loadNewFeatureIds(*pool_.slaveTransaction(), lastTxnId_);
    std::sort(newFeatureIds.begin(), newFeatureIds.end());
    common::forEachBatch(newFeatureIds, batchSize, [&](auto first, auto last) {
        auto features = db::FeatureGateway{*pool_.slaveTransaction()}.load(
            db::table::Feature::id.in({first, last}) &&
            (db::table::Feature::isPublished.is(true) ||
             db::table::Feature::shouldBePublished.is(true)));
        auto featureIds = toFeatureIds(features);
        auto objects = db::ObjectInPhotoGateway{*pool_.slaveTransaction()}.load(
            db::table::ObjectInPhotoTable::featureId.in(std::move(featureIds)));
        auto objectsMap = toObjectsInPhotoMap(std::move(objects));
        fn(std::move(features), std::move(objectsMap), Origin::Db);
    });
    return {newFeatureIds.begin(), newFeatureIds.end()};
}

void FbLoader::loadOldPublishedFeatures(size_t batchSize,
                                        const FeaturesFn& fn,
                                        const db::TIdSet& newFeatureIds) const
{
    auto features = db::Features{};
    auto objectsMap = ObjectsInPhotoMap{};
    auto readers = {std::cref(defaultFeatures_), std::cref(secretFeatures_)};
    for (const auto& reader : readers) {
        for (size_t i = 0; i < reader.get().featuresNumber(); ++i) {
            auto feature = reader.get().feature(i);
            if (newFeatureIds.contains(feature.id())) {
                continue;
            }
            auto objects = reader.get().objectsInPhoto(i);
            if (!objects.empty()) {
                objectsMap.insert({feature.id(), std::move(objects)});
            }
            features.push_back(std::move(feature));
            if (features.size() >= batchSize) {
                fn(std::move(features), std::move(objectsMap), Origin::Fb);
                features.clear();    // unspecified state to empty state
                objectsMap.clear();  // unspecified state to empty state
            }
        }
    }
    if (!features.empty()) {
        fn(std::move(features), std::move(objectsMap), Origin::Fb);
    }
}

void FbLoader::forEachPublishedFeatures(size_t batchSize, FeaturesFn fn) const
{
    auto newFeatureIds = loadNewPublishedFeatures(batchSize, fn);
    loadOldPublishedFeatures(batchSize, fn, newFeatureIds);
}

void FbLoader::forEachWalkObjects(size_t batchSize, WalkObjectsFn fn) const
{
    export_gen::forEachWalkObjects(pool_, batchSize, fn);
}

std::string_view FbLoader::fbVersion() const
{
    return defaultFeatures_.version();
}

std::unique_ptr<ILoader> makeLoader(
    pgpool3::Pool& pool,
    const std::string& mrcDatasetDir,
    const std::string& mrcFeaturesSecretDatasetDir,
    db::TId lastTxnId,
    EMappingMode mappingMode)
{
    try {
        return std::make_unique<FbLoader>(pool,
                                          mrcDatasetDir,
                                          mrcFeaturesSecretDatasetDir,
                                          lastTxnId,
                                          mappingMode);
    }
    catch (const std::exception& e) {
        WARN() << e.what();
    }
    return std::make_unique<DbLoader>(pool);
}

}  // namespace maps::mrc::export_gen
