#include "complex_objects_loader.h"
#include "filter_helpers.h"
#include <maps/wikimap/mapspro/services/editor/src/relations_manager.h>
#include <maps/wikimap/mapspro/services/editor/src/objects/category_traits.h>
#include <maps/wikimap/mapspro/services/editor/src/objects/complex_object.h>
#include <maps/wikimap/mapspro/services/editor/src/revision_meta/common.h>
#include <maps/wikimap/mapspro/services/editor/src/configs/config.h>
#include <maps/wikimap/mapspro/services/editor/src/configs/categories_strings.h>

#include <yandex/maps/wiki/common/batch.h>
#include <yandex/maps/wiki/common/string_utils.h>
#include <yandex/maps/wiki/revision/filters.h>
#include <yandex/maps/wiki/revision/historical_snapshot.h>

#include <yandex/maps/wiki/threadutils/executor.h>
#include <yandex/maps/wiki/threadutils/thread_pools.h>

namespace rf = maps::wiki::revision::filters;

namespace maps {
namespace wiki {

namespace {

const size_t COMMITS_BATCH_COUNT = 500;
const size_t OBJECTS_BATCH_COUNT = 500;
const size_t MINIMUM_THREAD_BATCH = 100;

const std::string& extractMasterCategory(const StringMap& attributes)
{
    auto it = attributes.find(ATTR_REL_MASTER);
    if (it != attributes.end()) {
        return it->second;
    }
    return s_emptyString;
}

bool isTopLevelMasterCategory(const Category& category)
{
    if (isIndirectGeomMasterCategory(category)) {
        return true;
    }

    if (!isContourCategory(category) && isDirectGeomMasterCategory(category)) {
        return true;
    }
    return false;
}

}

ComplexObjectsLoader::ComplexObjectsLoader(
        const BBoxContext& context,
        const CategoriesHolder& categoriesHolder,
        size_t threadPoolSize)
    : context_(context)
    , categoriesHolder_(categoriesHolder)
    , branchContext_(context_.getReadContext())
    , gateway_(branchContext_.txnCore(), branchContext_.branch)
    , threadPoolSize_(threadPoolSize)
    , failed_(false)
{
}

std::map<std::string, TOIds> ComplexObjectsLoader::topLevelComplexObjectIds()
{
    return std::move(topLevelComplexObjectIds_);
}

void ComplexObjectsLoader::run()
{
    std::map<TCommitId, TOIds> commitToObjects;

    loadHistoryGeometryObjects(commitToObjects);
    loadHistoryRelations(commitToObjects);

    if (!context_.checkOk(branchContext_.txnCore())) {
        return;
    }

    loadMasterRelationsByCommits(commitToObjects);

    commitToObjects.clear(); //free memory

    if (context_.customCommitIds().empty()) {
        return;
    }

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

    if (!context_.checkOk(branchContext_.txnCore())) {
        return;
    }

    INFO() << "Start loading commits with relations that will be added in commits range";

    auto newCommitToObjects = loadNewRelationsForObjects(loadedObjectIds_);
    if (newCommitToObjects.empty()) {
        return;
    }

    loadMasterRelationsByCommits(newCommitToObjects);

    TOIds newObjectIds;
    for (const auto& pair : newCommitToObjects) {
        const auto& objectIds = pair.second;
        newObjectIds.insert(objectIds.begin(), objectIds.end());
    }

    loadNewRelationsForObjects(newObjectIds);
}

bool ComplexObjectsLoader::checkRelationToTopLevelObject(const revision::ObjectRevision& rev)
{
    const auto& attributes = *rev.data().attributes;

    const auto& masterCategoryId = extractMasterCategory(attributes);
    if (!cfg()->editor()->categories().defined(masterCategoryId)) {
        return true;
    }
    const auto& masterCategory = cfg()->editor()->categories()[masterCategoryId];

    if (masterCategory.cacheGeomPartsThreshold()
            && isTopLevelMasterCategory(masterCategory)) {
        std::lock_guard<std::mutex> lock(topLevelComplexObjectsMutex_);
        const auto& relationData = *rev.data().relationData;
        topLevelComplexObjectIds_[masterCategoryId].insert(relationData.masterObjectId());
        return true;
    }
    return false;
}

void ComplexObjectsLoader::loadHistoryGeometryObjects(std::map<TCommitId, TOIds>& commitToObjects)
{
    INFO() << "Start loading commits with geometry objects";

    rf::ProxyFilterExpr geometryFilter = createObjectsFilter(categoriesHolder_.geoCategoryIds())
        && rf::Geom::defined()
        && rf::ObjRevAttr::isNotRelation();

    size_t counter = 0;

    if (context_.customCommitIds().empty()) {
        auto snapshot = gateway_.historicalSnapshot(context_.startCommitId(), context_.endCommitId());
        for (const auto& revId : snapshot.revisionIdsByFilter(geometryFilter)) {
            commitToObjects[revId.commitId()].insert(revId.objectId());
            counter++;
        }
    } else {
        common::applyBatchOp<TCommitIds>(
            context_.customCommitIds(),
            COMMITS_BATCH_COUNT,
            [&, this](const TCommitIds& batch) {
                auto filter = geometryFilter && rf::ObjRevAttr::commitId().in(batch);

                auto minCommitId = *batch.begin();
                auto maxCommitId = *batch.rbegin();

                auto snapshot = gateway_.historicalSnapshot(minCommitId, maxCommitId);
                for (const auto& revId : snapshot.revisionIdsByFilter(filter)) {
                    commitToObjects[revId.commitId()].insert(revId.objectId());
                    counter++;
                }
            });
    }

    //TODO: check for attribute only changes

    INFO() << "Loaded " << counter << " geometry object ids";
}

void ComplexObjectsLoader::loadHistoryRelations(std::map<TCommitId, TOIds>& commitToObjects)
{
    INFO() << "Start loading commits with relations";

    auto relationsFilter = createRelationsFilter(categoriesHolder_.complexCategoryIds());

    size_t counter = 0;

    if (context_.customCommitIds().empty()) {
        auto snapshot = gateway_.historicalSnapshot(context_.startCommitId(), context_.endCommitId());
        for (const auto& rev : snapshot.relationsByFilter(relationsFilter)) {
            if (!checkRelationToTopLevelObject(rev)) {
                auto commitId = rev.id().commitId();
                const auto& relationData = *rev.data().relationData;
                commitToObjects[commitId].insert(relationData.masterObjectId());
            }
            counter++;
        }
    } else {
        common::applyBatchOp<TCommitIds>(
            context_.customCommitIds(),
            COMMITS_BATCH_COUNT,
            [&, this](const TCommitIds& batch) {
                auto filter = relationsFilter && rf::ObjRevAttr::commitId().in(batch);

                auto minCommitId = *batch.begin();
                auto maxCommitId = *batch.rbegin();

                auto snapshot = gateway_.historicalSnapshot(minCommitId, maxCommitId);
                for (const auto& rev : snapshot.relationsByFilter(filter)) {
                    if (!checkRelationToTopLevelObject(rev)) {
                        auto commitId = rev.id().commitId();
                        const auto& relationData = *rev.data().relationData;
                        commitToObjects[commitId].insert(relationData.masterObjectId());
                    }
                    counter++;
                }
            });
    }

    INFO() << "Loaded " << counter << " relations";
}

template<typename Container, typename Func>
void ComplexObjectsLoader::doActionInThreadPool(const Container& data, size_t minimumBatchSize, Func func)
{
    if (threadPoolSize_ <= 1 || data.size() <= minimumBatchSize) {
        func(branchContext_.txnCore(), data);
        return;
    }

    Executor executor;

    size_t threadBatchSize = data.size() / threadPoolSize_ + 1;

    DEBUG() << "ComplexObjectsLoader: batch size per thread " << threadBatchSize;


    Container threadBatchData;
    auto task = [this, threadBatchData, func]() {
        try {
            if (!failed_) {
                auto branchContext = context_.getReadContext();
                func(branchContext.txnCore(), threadBatchData);
            }
            return;
        } catch (const maps::Exception& e) {
            ERROR() << "ComplexObjectsLoader: " << e;
        } catch (const std::exception& e) {
            ERROR() << "ComplexObjectsLoader: " << e.what();
        } catch (...) {
            ERROR() << "ComplexObjectsLoader: unknown error";
        }
        failed_ = true;
    };

    for (const auto& element : data) {
        threadBatchData.insert(element);

        if (threadBatchData.size() >= threadBatchSize) {
            executor.addTask(task);
            threadBatchData.clear();
        }
    }
    if (!threadBatchData.empty()) {
        executor.addTask(task);
    }

    auto poolHolder = cfg()->threadPools().pool(threadPoolSize_);
    executor.executeAllInThreads(*poolHolder);

    ASSERT(!failed_);
}

void ComplexObjectsLoader::loadMasterRelationsByCommits(const std::map<TCommitId, TOIds>& commitToObjects)
{
    INFO() << "Start loading objects from " << commitToObjects.size() << " commits";

    doActionInThreadPool(
        commitToObjects,
        MINIMUM_THREAD_BATCH,
        [&](pqxx::transaction_base& txn, const std::map<TCommitId, TOIds>& threadBatchCommitToObjects) {
            TOIds loadedObjectIdsInThread;

            for (const auto& it : threadBatchCommitToObjects) {
                auto commitId = it.first;
                const auto& objectIds = it.second;

                TOIds objectIdsToLoad;
                std::set_difference(
                    objectIds.begin(), objectIds.end(),
                    loadedObjectIdsInThread.begin(), loadedObjectIdsInThread.end(),
                    std::inserter(objectIdsToLoad, objectIdsToLoad.begin()));

                if (objectIdsToLoad.empty()) {
                    continue;
                }
                loadedObjectIdsInThread.insert(objectIdsToLoad.begin(), objectIdsToLoad.end());

                DEBUG() << "Load commit " << commitId << " object counts " << objectIds.size();
                auto newObjectIdsToLoad = loadMasterRelationsOnCommitInThread(txn, commitId, objectIds);

                if (newObjectIdsToLoad.empty()) {
                    continue;
                }
                loadedObjectIdsInThread.insert(newObjectIdsToLoad.begin(), newObjectIdsToLoad.end());

                DEBUG() << "Load commit " << commitId << " more object count " << newObjectIdsToLoad.size();
                loadMasterRelationsOnCommitInThread(txn, commitId, newObjectIdsToLoad);
            }

            std::lock_guard<std::mutex> lock(loadedObjectsIdsMutex_);
            loadedObjectIds_.insert(loadedObjectIdsInThread.begin(), loadedObjectIdsInThread.end());
        });
}

TOIds ComplexObjectsLoader::loadMasterRelationsOnCommitInThread(
    pqxx::transaction_base& txn,
    TCommitId commitId,
    const TOIds& slaveObjectIdsToLoad)
{
    if (slaveObjectIdsToLoad.empty()) {
        return {};
    }

    revision::RevisionsGateway gateway(txn, branchContext_.branch);
    auto snapshot = gateway.snapshot(commitId);

    TOIds newObjectIdsToLoad;

    common::applyBatchOp<TOIds>(
        slaveObjectIdsToLoad,
        OBJECTS_BATCH_COUNT,
        [&, this](const TOIds& batch) {
            auto relationsFilter = createRelationsFilter(categoriesHolder_.complexCategoryIds())
                && rf::ObjRevAttr::slaveObjectId().in(batch)
                && rf::ObjRevAttr::isNotDeleted();

            for (const auto& rev : snapshot.relationsByFilter(relationsFilter)) {
                if (!checkRelationToTopLevelObject(rev)) {
                    const auto& relationData = *rev.data().relationData;
                    newObjectIdsToLoad.insert(relationData.masterObjectId());
                }
            }
        });

    return newObjectIdsToLoad;
}

std::map<TCommitId, TOIds> ComplexObjectsLoader::loadNewRelationsForObjects(const TOIds& objectIds)
{
    INFO() << "Start loading new relations for " << objectIds.size() << " objects";

    std::mutex commitToObjectsMutex;
    std::map<TCommitId, TOIds> commitToObjects;

    doActionInThreadPool(
        objectIds,
        MINIMUM_THREAD_BATCH,
        [&](pqxx::transaction_base& txn, const TOIds& threadBatchObjectIds) {
            auto batchCommitToObjects = loadNewRelationsForObjectsInThread(txn, threadBatchObjectIds);

            std::lock_guard<std::mutex> lock(commitToObjectsMutex);

            for (const auto& pair : batchCommitToObjects) {
                auto commitId = pair.first;
                const auto& objectIds = pair.second;
                commitToObjects[commitId].insert(objectIds.begin(), objectIds.end());
            }
        });

    return commitToObjects;
}

std::map<TCommitId, TOIds> ComplexObjectsLoader::loadNewRelationsForObjectsInThread(
    pqxx::transaction_base& txn,
    const TOIds& objectIds)
{
    std::map<TCommitId, TOIds> commitToObjects;

    revision::RevisionsGateway gateway(txn, branchContext_.branch);

    auto minCommitId = *context_.customCommitIds().begin();
    auto maxCommitId = context_.endCommitId();
    auto snapshot = gateway.historicalSnapshot(minCommitId, maxCommitId);

    common::applyBatchOp<TOIds>(
        objectIds,
        OBJECTS_BATCH_COUNT,
        [&, this](const TOIds& batch) {
            auto filter = createRelationsFilter(categoriesHolder_.complexCategoryIds())
                && rf::ObjRevAttr::slaveObjectId().in(batch);

            for (auto&& rev : snapshot.relationsByFilter(filter)) {
                const auto& relationData = *rev.data().relationData;
                auto masterId = relationData.masterObjectId();
                if (!loadedObjectIds_.count(masterId)
                        && !checkRelationToTopLevelObject(rev)) {
                    auto commitId = rev.id().commitId();
                    commitToObjects[commitId].insert(masterId);
                }
            }
        });

    return commitToObjects;
}

} // namespace wiki
} // namespace maps
