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

#include <maps/libs/log8/include/log8.h>
#include <maps/libs/common/include/exception.h>
#include <yandex/maps/wiki/common/batch.h>
#include <yandex/maps/wiki/common/retry_duration.h>
#include <yandex/maps/wiki/common/revision_utils.h>
#include <yandex/maps/wiki/revision/exception.h>
#include <yandex/maps/wiki/revision/revisionsgateway.h>
#include <yandex/maps/wiki/revision/branch_manager.h>

namespace mwrev = maps::wiki::revision;
namespace rf = maps::wiki::revision::filters;
namespace mwcommon = maps::wiki::common;

namespace maps::wiki::validator {

namespace {

const size_t RELATION_LOAD_BATCH_SIZE = 10000;

struct RelationData {
    mwrev::DBID masterObjectId;
    mwrev::DBID slaveObjectId;
    decltype(Relation::otherCategoryId) masterCategoryId;
    decltype(Relation::otherCategoryId) slaveCategoryId;
    decltype(Relation::role) role;
    decltype(Relation::seqNum) seqNum;
};


RelationData loadRelationData(const mwrev::ObjectRevision& relation)
try {
    RelationData result;

    const auto& data = relation.data();
    REQUIRE(data.relationData,
            "Relation object revision: " << relation.id()
            << " doesn't contain relation data");
    result.masterObjectId = data.relationData->masterObjectId();
    result.slaveObjectId = data.relationData->slaveObjectId();

    REQUIRE(data.attributes,
            "Relation object revision: " << relation.id() << " has no attributes");

    auto checkedAttrValue = [&](const std::string& name)
    {
        auto nameValuePair = data.attributes->find(name);
        REQUIRE(nameValuePair != data.attributes->end(),
            "Relation object revision: " << relation.id()
            << " doesn't contain " << name << " attribute");
        return nameValuePair->second;
    };

    result.masterCategoryId = checkedAttrValue(REL_MASTER_ATTR);
    result.slaveCategoryId = checkedAttrValue(REL_SLAVE_ATTR);
    result.role = checkedAttrValue(REL_ROLE_ATTR);

    auto seqNum = data.attributes->find(REL_SEQ_NUM_ATTR);
    if (seqNum != data.attributes->end()) {
        try {
            result.seqNum = std::stoul(seqNum->second);
        } catch (std::logic_error&) {
            throw maps::RuntimeError()
                    << "Relation object revision: " << relation.id()
                    << " contains malformed " << REL_SEQ_NUM_ATTR
                    << " attribute";
        }
    }

    return result;
} catch (const std::exception& ex) {
    throw common::RetryDurationCancel() << ex.what();
}

revision::Branch loadBranch(pgpool3::Pool& pgPool, revision::DBID branchId)
{
    return common::retryDuration([&] {
        auto txn = pgPool.masterReadOnlyTransaction();
        revision::BranchManager branchManager(*txn);
        try {
            return branchManager.load(branchId);
        } catch (const revision::BranchNotExistsException& ex) {
            throw common::RetryDurationCancel() << ex.what();
        }
    });
}

} // namespace

DBGateway::DBGateway(
        pgpool3::Pool& pgPool,
        revision::DBID branchId,
        revision::DBID commitId)
    : pgPool_(pgPool)
    , branch_(loadBranch(pgPool_, branchId))
    , commitId_(commitId)
{ }

size_t DBGateway::maxConnections() const
{
    return pgPool_.state().constants.slaveMaxSize;
}

pgpool3::TransactionHandle DBGateway::getTransaction() const
{
    return mwcommon::getReadTransactionForCommit(
        pgPool_, branch_.id(), commitId_,
        [] (const std::string& s) { INFO() << s; });
}

revision::Snapshot DBGateway::snapshot(pqxx::transaction_base& txn) const
{
    mwrev::RevisionsGateway gateway(txn, branch_);
    return gateway.stableSnapshot(commitId_);
}

RevisionIds DBGateway::revisionIdsByFilter(
        const mwrev::filters::FilterExpr& filter,
        pqxx::transaction_base& txn) const
{
    return snapshot(txn).revisionIdsByFilter(filter);
}

RevisionIds DBGateway::revisionIdsByObjectIds(
        const revision::ConstRange<TId>& ids,
        pqxx::transaction_base& txn) const
{
    return snapshot(txn).objectRevisionIds(ids);
}

RevisionIds DBGateway::revisionIdsByObjectIds(
        const revision::ConstRange<TId>& ids,
        const revision::filters::FilterExpr& filter,
        pqxx::transaction_base& txn) const
{
    return snapshot(txn).revisionIdsByFilter(ids, filter);
}

Revisions DBGateway::revisionsByRevisionIds(
        const RevisionIds& ids,
        bool hasGeom,
        pqxx::transaction_base& txn) const
{
    mwrev::RevisionsGateway gateway(txn, branch_);
    auto partitionHint =
        rf::ObjRevAttr::isNotRelation()
        && (hasGeom
            ? rf::ProxyFilterExpr(rf::Geom::defined())
            : rf::ProxyFilterExpr(!rf::Geom::defined()));
    return gateway.reader().loadRevisions(ids, partitionHint);
}

RelationsMap DBGateway::mastersOf(
        const std::vector<TId>& slaves,
        pqxx::transaction_base& txn) const
{
    RelationsMap ret;

    auto sn = snapshot(txn);

    common::applyBatchOp<std::vector<TId>>(
        slaves,
        RELATION_LOAD_BATCH_SIZE,
        [&](const std::vector<TId>& batch) {
            for (const auto& revision : sn.loadMasterRelations(batch)) {
                auto data = loadRelationData(revision);
                ret[data.slaveObjectId].emplace_back(
                    data.masterObjectId,
                    data.masterCategoryId,
                    data.role,
                    data.seqNum
                );
            }
        });

    return ret;
}


RelationsMap DBGateway::slavesOf(
        const std::vector<TId>& masters,
        pqxx::transaction_base& txn) const
{
    RelationsMap ret;

    auto sn = snapshot(txn);

    common::applyBatchOp<std::vector<TId>>(
        masters,
        RELATION_LOAD_BATCH_SIZE,
        [&](const std::vector<TId>& batch) {
            for (const auto& revision : sn.loadSlaveRelations(batch)) {
                auto data = loadRelationData(revision);
                ret[data.masterObjectId].emplace_back(
                    data.slaveObjectId,
                    data.slaveCategoryId,
                    data.role,
                    data.seqNum
                );
            }
        });

    return ret;
}

ObjectIdSet DBGateway::loadObjectIdsByRelatedIds(
    pqxx::transaction_base& txn,
    const TCategoryId& categoryId,
    const std::vector<TId>& relatedIds,
    RelationType relationType,
    const std::set<std::string>& roleIds) const
{
    ASSERT(!roleIds.empty());

    ObjectIdSet objectIdsSet;

    auto sn = snapshot(txn);

    const auto& categoryAttr = relationType == RelationType::Slave
        ? REL_MASTER_ATTR
        : REL_SLAVE_ATTR;

    auto filter =
        rf::ObjRevAttr::isNotDeleted() &&
        rf::Attr(REL_ROLE_ATTR).in(roleIds) &&
        rf::Attr(categoryAttr) == categoryId;

    common::applyBatchOp<std::vector<TId>>(
        relatedIds,
        RELATION_LOAD_BATCH_SIZE,
        [&](const std::vector<TId>& batch) {
            if (relationType == RelationType::Slave) {
                for (const auto& relation : sn.loadMasterRelations(batch, filter)) {
                    const auto& data = relation.data();
                    ASSERT(data.relationData);
                    objectIdsSet.insert(data.relationData->masterObjectId());
                }
            } else {
                for (const auto& relation : sn.loadSlaveRelations(batch, filter)) {
                    const auto& data = relation.data();
                    ASSERT(data.relationData);
                    objectIdsSet.insert(data.relationData->slaveObjectId());
                }
            }
        });
    return objectIdsSet;
}

} // namespace maps::wiki::validator
