#include "exporter.h"
#include "jsonhelper.h"

#include <yandex/maps/wiki/revisionapi/export_params.h>

#include <yandex/maps/wiki/revision/common.h>
#include <yandex/maps/wiki/revision/objectrevision.h>
#include <yandex/maps/wiki/revision/revisionid.h>
#include <yandex/maps/wiki/revision/revisionsgateway.h>
#include <yandex/maps/wiki/common/constants.h>
#include <yandex/maps/wiki/common/revision_utils.h>
#include <maps/libs/geolib/include/conversion.h>

#include <boost/algorithm/string/join.hpp>

#include <algorithm>
#include <functional>
#include <iterator>
#include <unordered_map>
#include <unordered_set>
#include <utility>

namespace filters = revision::filters;

namespace maps::wiki::revisionapi {

namespace {

const size_t DB_READ_BATCH_SIZE = 10000;
const std::string REL_ROLE_ATTR = "rel:role";
const std::string REL_MASTER_ATTR = "rel:master";
const std::string REL_SLAVE_ATTR = "rel:slave";

struct ObjectWithRelations {
    explicit ObjectWithRelations(const revision::ObjectRevision& object)
        : object(object)
    {
    }

    void json(json::ObjectBuilder builder) const
    try {
        const auto& objectData = object.data();

        if (objectData.attributes) {
            builder[JSON_FIELD_ATTRIBUTES] = [&](json::ObjectBuilder builder) {
                attributes2json(builder, *objectData.attributes);
            };
        }
        if (objectData.description) {
            builder[JSON_FIELD_DESCRIPTION] = *objectData.description;
        }
        if (objectData.geometry) {
            auto mercatorGeometry =
                geolib3::WKB::read<geolib3::SimpleGeometryVariant>(*objectData.geometry);
            auto geoGeometry = geolib3::convertMercatorToGeodetic(mercatorGeometry);

            builder[JSON_FIELD_GEOMETRY] = geolib3::geojson(geoGeometry);
        }

        if (masterToSlaveRelations.empty() && slaveToMasterRelations.empty()) {
            return;
        }
        builder[JSON_FIELD_RELATIONS] = [&](json::ArrayBuilder builder) {
            for (const auto& relation: masterToSlaveRelations) {
                builder << [&](json::ObjectBuilder builder) {
                    builder[JSON_FIELD_SLAVE] = std::to_string(
                        relation.data().relationData->slaveObjectId());
                    builder[JSON_FIELD_ATTRIBUTES] = [&](json::ObjectBuilder builder) {
                        attributes2json(builder, *relation.data().attributes);
                    };
                };
            }
            for (const auto& relation: slaveToMasterRelations) {
                builder << [&](json::ObjectBuilder builder) {
                    builder[JSON_FIELD_MASTER] = std::to_string(
                        relation.data().relationData->masterObjectId());
                    builder[JSON_FIELD_ATTRIBUTES] = [&](json::ObjectBuilder builder) {
                        attributes2json(builder, *relation.data().attributes);
                    };
                };
            }
        };
    } catch (const std::exception& ex) {
        throw DataError() << "Can not convert object to json " << object.id() << " : " << ex.what();
    }

    revision::ObjectRevision object;
    std::vector<revision::ObjectRevision> masterToSlaveRelations;
    std::vector<revision::ObjectRevision> slaveToMasterRelations;
};

struct Chunk {
    std::vector<ObjectWithRelations> objects;
};

bool
allowedObjectCategory(
    const std::set<std::string>& categoriesForFilter,
    const revision::ObjectRevision& r)
{
    const auto& attributes = r.data().attributes;
    if (!attributes) {
        return false;
    }
    if (categoriesForFilter.empty()) {
        return true;
    }

    for (const auto& pair : *attributes) {
        if (categoriesForFilter.contains(pair.first)) {
            return true;
        }
    }
    return false;
}

bool
allowedRelation(
    const std::set<std::string>& categories,
    const revision::ObjectRevision& r)
{
    const auto& attributes = r.data().attributes;
    if (!attributes) {
        return false;
    }

    auto itMaster = attributes->find(REL_MASTER_ATTR);
    REQUIRE(itMaster != attributes->end(),
            "master category not found, " << r.id());

    auto itSlave = attributes->find(REL_SLAVE_ATTR);
    REQUIRE(itSlave != attributes->end(),
            "slave category not found, " << r.id());

    if (categories.empty()) {
        return true;
    }

    return categories.contains(itMaster->second) &&
           categories.contains(itSlave->second);
}


class ExportHelper
{
public:
    ExportHelper(
            const revision::Snapshot& snapshot,
            const ExportParams& params,
            GetStreamForChunkFunc getStreamForChunkFunc,
            const revision::DBID lastObjectId)
        : snapshot_(snapshot)
        , params_(params)
        , getStreamForChunkFunc_(std::move(getStreamForChunkFunc))
        , lastObjectId_(lastObjectId)
        , currentWriteChunkId_(0)
    {}

    void perform(
        const revision::RevisionIds& revIds,
        const FilterPtr& filter = FilterPtr());

private:
    std::unique_ptr<Chunk> readChunk(
        const revision::RevisionIds& objectRevisionIds,
        const std::unordered_set<revision::DBID>& exportOidsSet,
        const FilterPtr& filter) const;

    const revision::Snapshot& snapshot_;
    const ExportParams& params_;
    const GetStreamForChunkFunc getStreamForChunkFunc_;
    const revision::DBID lastObjectId_;
    size_t currentWriteChunkId_;
};

std::unique_ptr<Chunk>
ExportHelper::readChunk(
    const revision::RevisionIds& objectRevisionIds,
    const std::unordered_set<revision::DBID>& exportOidsSet,
    const FilterPtr& filter) const
{
    auto chunk = std::make_unique<Chunk>();
    chunk->objects.reserve(objectRevisionIds.size());
    std::unordered_map<revision::DBID, ObjectWithRelations*> dbidToObject;
    std::vector<revision::DBID> objectDbIds;

    filters::ProxyFilterExpr objectsFilter(filters::ObjRevAttr::isNotRelation());
    if (filter) {
        objectsFilter &= filters::ProxyFilterExpr(filter);
    }

    const auto revisions =
        snapshot_.reader().loadRevisions(objectRevisionIds, objectsFilter);

    const auto& categoriesForFilter = params_.categoriesForFilter();
    for (const auto& r: revisions) {
        if (!allowedObjectCategory(categoriesForFilter, r)) {
            continue;
        }
        auto dbid = r.id().objectId();
        chunk->objects.emplace_back(r);
        dbidToObject.insert({dbid, &chunk->objects.back()});
        objectDbIds.push_back(dbid);
    }
    if (objectDbIds.empty()) {
        return chunk;
    }

    const auto flags = params_.flags();
    const auto& allowedRelativesCategories = params_.allowedRelativesCategories();

    auto processMasterToSlaveRelation = [&] (const revision::ObjectRevision& r)
    {
        auto it = dbidToObject.find(r.data().relationData->masterObjectId());
        REQUIRE(it != dbidToObject.end(),
                "master object for slave not found, " << r.id());
        if (!allowedRelation(allowedRelativesCategories, r)) {
            return;
        }
        if (!(flags & RelationsExportFlags::SkipDangling)
                || exportOidsSet.contains(r.data().relationData->slaveObjectId())) {
            it->second->masterToSlaveRelations.push_back(r);
        }
    };

    auto processSlaveToMasterRelation = [&] (const revision::ObjectRevision& r)
    {
        auto it = dbidToObject.find(r.data().relationData->slaveObjectId());
        REQUIRE(it != dbidToObject.end(),
                "slave object for master not found, " << r.id());
        if (!allowedRelation(allowedRelativesCategories, r)) {
            return;
        }
        if (!(flags & RelationsExportFlags::SkipDangling)
                || exportOidsSet.contains(r.data().relationData->masterObjectId())) {
            it->second->slaveToMasterRelations.push_back(r);
        }
    };

    const auto relationsFilter = filters::ObjRevAttr::isNotDeleted();
    const auto& invertDirectionRoles = params_.invertDirectionRoles();

    if (!!(flags & RelationsExportFlags::MasterToSlave)) {
        filters::ProxyFilterExpr filter(relationsFilter);
        filter &= filters::ObjRevAttr::masterObjectId().in(objectDbIds);
        if (!invertDirectionRoles.empty()) {
            filter &= !filters::Attr(REL_ROLE_ATTR).in(invertDirectionRoles);
        }
        for (const auto& r: snapshot_.relationsByFilter(filter)) {
            processMasterToSlaveRelation(r);
        }

        if (!invertDirectionRoles.empty()) {
            filters::ProxyFilterExpr filter(relationsFilter);
            filter &= filters::ObjRevAttr::slaveObjectId().in(objectDbIds);
            filter &= filters::Attr(REL_ROLE_ATTR).in(invertDirectionRoles);

            for (const auto& r: snapshot_.relationsByFilter(filter)) {
                processSlaveToMasterRelation(r);
            }
        }
    }

    if (!!(flags & RelationsExportFlags::SlaveToMaster)) {
        filters::ProxyFilterExpr filter(relationsFilter);
        filter &= filters::ObjRevAttr::slaveObjectId().in(objectDbIds);
        if (!invertDirectionRoles.empty()) {
            filter &= !filters::Attr(REL_ROLE_ATTR).in(invertDirectionRoles);
        }
        for (const auto& r: snapshot_.relationsByFilter(filter)) {
            processSlaveToMasterRelation(r);
        }

        if (!invertDirectionRoles.empty()) {
            filters::ProxyFilterExpr filter(relationsFilter);
            filter &= filters::ObjRevAttr::masterObjectId().in(objectDbIds);
            filter &= filters::Attr(REL_ROLE_ATTR).in(invertDirectionRoles);

            for (const auto& r: snapshot_.relationsByFilter(filter)) {
                processMasterToSlaveRelation(r);
            }
        }
    }

    return chunk;
};

void
ExportHelper::perform(
    const revision::RevisionIds& revIds,
    const FilterPtr& filter)
{
    std::unordered_set<revision::DBID> exportOidsSet;
    if (!!(params_.flags() & RelationsExportFlags::SkipDangling)) {
        for (const auto& revId : revIds) {
            exportOidsSet.insert(revId.objectId());
        }
    }

    const auto nextFreeObjectId = std::to_string(lastObjectId_ + 1);
    const auto relationsExportMode =
        boost::lexical_cast<std::string>(params_.flags());
    const auto invertDirectionsForRoles =
        boost::join(params_.invertDirectionRoles(), ",");

    auto readIter = revIds.begin();
    std::unique_ptr<Chunk> chunk(new Chunk());
    auto writeIter = chunk->objects.begin();

    auto updateReadChunk = [&]()
    {
        while (writeIter == chunk->objects.end() && readIter != revIds.end()) {
            const size_t totalReadLeft = std::distance(readIter, revIds.end());
            const size_t readChunkSize = std::min(DB_READ_BATCH_SIZE, totalReadLeft);
            auto nextReadIter = readIter + readChunkSize;
            chunk = readChunk({readIter, nextReadIter}, exportOidsSet, filter);
            writeIter = chunk->objects.begin();
            readIter = nextReadIter;
        }
        return writeIter != chunk->objects.end();
    };

    const size_t writeBatchSize = params_.writeBatchSize();
    auto createJsonChunk = [&]()
    {
        std::shared_ptr<std::ostream> os(getStreamForChunkFunc_(currentWriteChunkId_++));
        json::Builder builder(*os);
        builder.setDoublePrecision(common::JSON_DOUBLE_PRECISION);
        builder << [&](json::ObjectBuilder builder) {
            builder[JSON_FIELD_ATTRIBUTES] = [&](json::ObjectBuilder builder) {
                builder[JSON_FIELD_DESCRIPTION] = "exported from revisions";
                builder[JSON_FIELD_NEXT_FREE_OBJECT_ID] = nextFreeObjectId;
                builder[JSON_FIELD_RELATIONS_EXPORT_MODE] = relationsExportMode;
                if (!invertDirectionsForRoles.empty()) {
                    builder[JSON_FIELD_INVERT_DIRECTIONS_FOR_ROLES] = invertDirectionsForRoles;
                }
            };

            builder[JSON_FIELD_OBJECTS] = [&](json::ObjectBuilder builder) {
                for (size_t batchSize = 0; updateReadChunk(); ) {
                    auto id = writeIter->object.id().objectId();
                    builder[std::to_string(id)] = *writeIter++;
                    if (writeBatchSize && ++batchSize >= writeBatchSize) {
                        break;
                    }
                }
            };
        };
    };

    do {
        if (updateReadChunk() || params_.emptyJsonPolicy() == EmptyJsonPolicy::Export) {
            createJsonChunk();
        }
    } while (updateReadChunk());
}


class ObjectRevisionIdsLoader
{
public:
    explicit ObjectRevisionIdsLoader(const revision::Snapshot& snapshot)
        : snapshot_(snapshot)
    {}

    revision::RevisionIds loadGeomRevisionIds(const FilterPtr& filter)
    {
        return snapshot_.revisionIdsByFilter(
            joinObjectsFilter(filter) && filters::Geom::defined());
    }

    revision::RevisionIds loadNoGeomRevisionIds(const FilterPtr& filter)
    {
        return snapshot_.revisionIdsByFilter(
            joinObjectsFilter(filter) && !filters::Geom::defined());
    }

    revision::RevisionIds loadRevisionIds(const FilterPtr& filter)
    {
        return joinRevisionIds(
            loadGeomRevisionIds(filter),
            loadNoGeomRevisionIds(filter));
    }

    revision::RevisionIds loadRevisionIds(
        const std::vector<revision::DBID>& objectIds)
    {
        revision::RevisionIds result;

        auto first = std::begin(objectIds);
        while (first != std::end(objectIds)) {
            auto last = first + std::min<size_t>(
                DB_READ_BATCH_SIZE, std::distance(first, std::end(objectIds)));

            auto curBatch = loadRevisionIds(
                std::make_shared<filters::ProxyFilterExpr>(
                    filters::ObjRevAttr::objectId().in({first, last})));

            std::move(
                std::begin(curBatch), std::end(curBatch),
                std::back_inserter(result));

            first = last;
        }
        return result;
    }

private:
    static revision::RevisionIds joinRevisionIds(
        revision::RevisionIds&& revIds1,
        revision::RevisionIds&& revIds2)
    {
        if (revIds1.size() < revIds2.size()) {
            revIds1.swap(revIds2);
        }
        if (!revIds2.empty()) {
            revIds1.insert(revIds1.end(), revIds2.begin(), revIds2.end());
        }
        return std::move(revIds1);
    }

    static filters::ProxyFilterExpr joinObjectsFilter(const FilterPtr& filter)
    {
        filters::ProxyFilterExpr filterAccum(filters::ObjRevAttr::isNotDeleted());
        filterAccum &= filters::ObjRevAttr::isNotRelation();
        if (filter) {
            filterAccum &= filters::ProxyFilterExpr(filter);
        }
        return filterAccum;
    }

    const revision::Snapshot& snapshot_;
};


FilterPtr makeCommonFilter(
    const ExportParams& params,
    const FilterPtr& filter)
{
    if (filter && params.filter()) {
        return std::make_shared<filters::BinaryFilterExpr>(
            filters::BinaryFilterExpr::Operation::And,
            filter,
            params.filter());
    }
    if (filter) {
        return filter;
    }
    return params.filter();
}

pgpool3::TransactionHandle getReadTransactionForCommit(
    pgpool3::Pool& pool,
    const ExportParams& params)
{
    return common::getReadTransactionForCommit(
        pool,
        params.branch().id(),
        params.commitId(),
        [] (const std::string&) {}
    );
}

} // namespace

Exporter::Exporter(pgpool3::Pool& pool)
    : pool_(pool)
{
}

void Exporter::exportToJsonByFilter(
    const ExportParams& params,
    GetStreamForChunkFunc getStreamForChunkFunc,
    const FilterPtr& filter)
{
    auto commonFilter = makeCommonFilter(params, filter);

    auto readTr = getReadTransactionForCommit(pool_, params);

    revision::RevisionsGateway gateway(*readTr, params.branch());
    gateway.reader().setDescriptionLoadingMode(revision::DescriptionLoadingMode::Load);

    auto snapshot = gateway.stableSnapshot(params.commitId());
    ObjectRevisionIdsLoader loader(snapshot);

    ExportHelper helper(
        snapshot, params, std::move(getStreamForChunkFunc), gateway.lastObjectId());

    if (!params.writeBatchSize()) {
        helper.perform(loader.loadRevisionIds(commonFilter));
        return;
    }

    auto geomFilter =
        std::make_shared<filters::GeomFilterExpr>(filters::Geom::defined());
    auto noGeomFilter =
        std::make_shared<filters::NegativeFilterExpr>(geomFilter);

    helper.perform(loader.loadGeomRevisionIds(commonFilter), geomFilter);
    helper.perform(loader.loadNoGeomRevisionIds(commonFilter), noGeomFilter);
}

void Exporter::exportToJsonByObjectIds(
    const ExportParams& params,
    GetStreamForChunkFunc getStreamForChunkFunc,
    const std::vector<revision::DBID>& objectIds)
{
    REQUIRE(!params.filter(),
        "Cannot set filter and object ids list simultaneously");

    auto readTr = getReadTransactionForCommit(pool_, params);

    revision::RevisionsGateway gateway(*readTr, params.branch());
    gateway.reader().setDescriptionLoadingMode(revision::DescriptionLoadingMode::Load);

    auto snapshot = gateway.stableSnapshot(params.commitId());

    ObjectRevisionIdsLoader loader(snapshot);

    ExportHelper helper(
        snapshot, params, std::move(getStreamForChunkFunc), gateway.lastObjectId());
    helper.perform(loader.loadRevisionIds(objectIds));
}

} // namespace maps::wiki::revisionapi
