#include "revisions_loaders.h"
#include "helper_loaders.h"

#include "data_source.h"
#include <maps/wikimap/mapspro/libs/validator/common/magic_strings.h>

#include <yandex/maps/wiki/validator/common.h>
#include <yandex/maps/wiki/common/retry_duration.h>
#include <yandex/maps/wiki/configs/editor/categories.h>
#include <yandex/maps/wiki/revision/filters.h>

#include <algorithm>
#include <vector>

namespace maps::wiki::validator {

namespace rf = revision::filters;

namespace {

rf::BinaryFilterExpr
filterForCategory(const TCategoryId& categoryId)
{
    return rf::Attr(CAT_PREFIX + categoryId).defined()
        && rf::ObjRevAttr::isNotDeleted()
        && rf::ObjRevAttr::isNotRelation()
        && (categoryGeomType(categoryId) != GeomType::None
            ? rf::ProxyFilterExpr(rf::Geom::defined())
            : rf::ProxyFilterExpr(!rf::Geom::defined()));
}

RevisionIds loadByRelatedIds(
    DataSource& dataSource,
    const TCategoryId& categoryId,
    const std::set<TCategoryId>& otherCategories,
    RelationType relationType,
    const std::set<std::string>& roleIds)
{
    ASSERT(!dataSource.aoi().empty() || dataSource.hasSelectedObjectIds());

    auto objectIdsSet = dataSource.selectedObjectIds(categoryId);

    for (const TCategoryId& otherCategory : otherCategories) {
        const ObjectsCollectionBase& otherCollection =
            dataSource.collection(otherCategory);
        otherCollection.checkLoaded();
        std::vector<TId> relatedIds = otherCollection.objectIds();

        auto objectIds = loadObjectIdsByRelatedIds(
            dataSource,
            categoryId,
            relatedIds,
            relationType,
            roleIds);
        objectIdsSet.insert(objectIds.begin(), objectIds.end());
    }

    return revisionIdsByObjectIds(dataSource, objectIdsSet);
}

RevisionIds loadRevisionIdsInsideAoi(
    DataSource& dataSource,
    const TCategoryId& categoryId)
{
    const auto& aoi = dataSource.aoi();
    ASSERT(!aoi.empty());
    ASSERT(!dataSource.hasSelectedObjectIds());

    const auto& bboxes = aoi.boundingBoxesBuffered();
    ASSERT(!bboxes.empty());

    rf::GeomFilterExpr::Operation op;
    switch (categoryGeomType(categoryId)) {
        case GeomType::Point:
            op = rf::GeomFilterExpr::Operation::IntersectsPoints;
            break;
        case GeomType::Polygon:
            op = rf::GeomFilterExpr::Operation::IntersectsPolygons;
            break;
        case GeomType::Polyline:
            op = rf::GeomFilterExpr::Operation::IntersectsLinestrings;
            break;
        default:
            throw LogicError() << "bad geom type for category: " << categoryId;
    }

    auto loadRevisionIds = [&](const auto& bbox) {
        rf::ProxyFilterExpr filterAccumulator(filterForCategory(categoryId));
        filterAccumulator &= rf::GeomFilterExpr(
            op,
            bbox.minX(), bbox.minY(), bbox.maxX(), bbox.maxY());

        return revisionIdsByFilter(dataSource, filterAccumulator);
    };

    if (bboxes.size() == 1) {
        return loadRevisionIds(bboxes.front());
    }

    auto bboxIt = bboxes.begin();
    auto result = loadRevisionIds(*bboxIt);

    ObjectIdSet objectIds;
    objectIds.reserve(result.size());
    for (const auto& revId : result) {
        objectIds.emplace(revId.objectId());
    }

    while (++bboxIt != bboxes.end()) {
        auto revIds = loadRevisionIds(*bboxIt);
        objectIds.reserve(objectIds.size() + revIds.size());
        result.reserve(result.size() + revIds.size());
        for (const auto& revId : revIds) {
            if (objectIds.insert(revId.objectId()).second) {
                result.push_back(revId);
            }
        }
    }
    return result;
}

RevisionIds loadAdmUnitRevisionIds(DataSource& dataSource)
{
    ASSERT(!dataSource.aoi().empty());
    ASSERT(!dataSource.hasSelectedObjectIds());

    std::vector<TId> slaveIds;

    const ObjectsCollectionBase& faces =
        dataSource.collection(categories::AD_FC::id());
    faces.checkLoaded();
    for (TId id : faces.objectIds()) {
        slaveIds.push_back(id);
    }
    const ObjectsCollectionBase& centers =
        dataSource.collection(categories::AD_CNT::id());
    centers.checkLoaded();
    for (const TId& id : centers.objectIds()) {
        slaveIds.push_back(id);
    }
    const ObjectsCollectionBase& substs =
        dataSource.collection(categories::AD_SUBST::id());
    substs.checkLoaded();
    for (const TId& id : substs.objectIds()) {
        slaveIds.push_back(id);
    }

    static const std::set<std::string> SLAVE_ROLE_IDS =
        {PART_ROLE, CENTER_ROLE, SUBSTITUTION_ROLE, AD_SUBST_CHILD_ROLE};

    auto objectIds = common::retryDuration([&] {
        auto txn = dataSource.getTransaction();

        ObjectIdSet objectIds;
        for (const auto& slaveToMasters : dataSource.dbGateway().mastersOf(slaveIds, *txn)) {
            for (const auto& relation : slaveToMasters.second) {
                if (SLAVE_ROLE_IDS.count(relation.role)) {
                    objectIds.insert(relation.other);
                }
            }
        }

        std::vector<TId> curBatch(objectIds.begin(), objectIds.end());
        std::vector<TId> nextBatch;
        while (!curBatch.empty()) {
            for (const auto& slaveToMasters : dataSource.dbGateway().mastersOf(curBatch, *txn)) {
                for (const auto& relation : slaveToMasters.second) {
                    if (relation.role == CHILD_ROLE) {
                        auto insertionResult = objectIds.insert(relation.other);
                        if (insertionResult.second) {
                            nextBatch.push_back(relation.other);
                        }
                    }
                }
            }
            curBatch.swap(nextBatch);
            nextBatch.clear();
        }
        return objectIds;
    });

    return revisionIdsByObjectIds(dataSource, objectIds);
}

} // namespace


RevisionIdsLoader createAllRevisionIdsLoader(const TCategoryId& categoryId)
{
    RevisionIdsLoader loader;

    loader.load = [categoryId](DataSource& dataSource) {
        return revisionIdsByFilter(dataSource, filterForCategory(categoryId));
    };

    return loader;
}

RevisionIdsLoader createByRelatedIdsLoader(
    const TCategoryId& categoryId,
    const std::set<TCategoryId>& otherCategories,
    const std::set<std::string>& roleIds,
    RelationType relationType)
{
    RevisionIdsLoader loader;

    loader.load =
        [categoryId, otherCategories, roleIds, relationType]
        (DataSource& dataSource) {
            return loadByRelatedIds(
                dataSource, categoryId, otherCategories, relationType, roleIds);
        };

    loader.dependencies = otherCategories;

    return loader;
}

RevisionIdsLoader createInsideAoiRevisionIdsLoader(const TCategoryId& categoryId)
{
    RevisionIdsLoader loader;

    loader.load = [categoryId](DataSource& dataSource) {
        return loadRevisionIdsInsideAoi(dataSource, categoryId);
    };

    return loader;
}

RevisionIdsLoader createAdmUnitLoader()
{
    RevisionIdsLoader loader;

    loader.load = [](DataSource& dataSource) {
        return loadAdmUnitRevisionIds(dataSource);
    };

    loader.dependencies = {
        categories::AD_FC::id(),
        categories::AD_CNT::id(),
        categories::AD_SUBST::id()
    };

    return loader;
}

} // namespace maps::wiki::validator
