#include "update_revision_attrs.h"

#include <yandex/maps/wiki/revision/revisionsgateway.h>
#include <yandex/maps/wiki/revision/branch_manager.h>
#include <yandex/maps/wiki/revision/filters.h>

#include <maps/libs/common/include/exception.h>
#include <maps/libs/log8/include/log8.h>

#include <pqxx>

#include <algorithm>

using namespace maps;
using namespace maps::wiki;

namespace {

const size_t ATTEMPTS = 5;
const std::string CATEGORY_PREFIX = "cat:";

class GlobalContext : public Options
{
public:
    GlobalContext(const std::string& connString, const Options& options)
        : Options(options)
        , connString_(connString)
        , connRead_(connString_)
        , connWrite_(connString_)
        , workRead_(connRead_)
        , branch_(revision::BranchManager(workRead_).load(options.branchId))
        , gw_(workRead_, branch_)
    {
        refreshDraftCommits();
    }

    void refreshDraftCommits()
    {
        if (!updateApprovedOnly) {
            return;
        }

        revision::DBID maxDraftCommitId = draftCommitIds_.empty()
            ? 0
            : *draftCommitIds_.rbegin();

        auto commits = revision::Commit::load(
            workRead_,
            revision::filters::CommitAttr::isDraft() &&
            revision::filters::CommitAttr::id() > maxDraftCommitId);
        for (const auto& commit : commits) {
            draftCommitIds_.insert(commit.id());
        }
    }

    revision::RevisionIds filterRevisionIds(revision::RevisionIds revIds)
    {
        revision::RevisionIds result;
        if (!updateApprovedOnly) {
            result.swap(revIds);
        } else {
            result.reserve(revIds.size());
            for (const auto& revId : revIds) {
                if (draftCommitIds_.count(revId.commitId())) {
                    WARN() << "Skip draft object revision: " << revId;
                    updateResult_.draftObjectIds.insert(revId.objectId());
                } else {
                    result.push_back(revId);
                }
            }
        }
        return result;
    }

    const revision::RevisionsGateway& gateway() const { return gw_; }

    void createCommit(
        const std::list<revision::RevisionsGateway::NewRevisionData>& newObjectsData,
        const revision::Attributes& commitAttributes,
        revision::UserID updatedBy)
    {
        if (dryRun || newObjectsData.empty()) {
            return;
        }

        pqxx::work work(connWrite_);
        revision::RevisionsGateway gw(work, branch_);
        auto commit = gw.createCommit(newObjectsData, updatedBy, commitAttributes);
        work.commit();
        INFO() << "Created commit " << commit.id();
        updateResult_.commitIds.push_back(commit.id());
        updateResult_.updateCount += newObjectsData.size();
    }

    UpdateResult& updateResult() { return updateResult_; }

private:
    std::string connString_;
    pqxx::connection connRead_;
    pqxx::connection connWrite_;
    pqxx::work workRead_;

    revision::Branch branch_;
    revision::RevisionsGateway gw_;
    revision::DBIDSet draftCommitIds_;

    UpdateResult updateResult_;
};

revision::RevisionIds
loadRevisionIds(
    GlobalContext& ctx,
    const std::vector<revision::DBID>& objectIds)
{
    const auto& gw = ctx.gateway();
    auto revIds = gw.snapshot(gw.headCommitId()).objectRevisionIds(objectIds);
    INFO() << revIds.size() << " revisions ids loaded";

    revision::DBIDSet loadedObjectIds;
    for (const auto& revId : revIds) {
        loadedObjectIds.insert(revId.objectId());
    }
    for (auto objectId : objectIds) {
        if (!loadedObjectIds.count(objectId)) {
            WARN() << "Non-existent object id: " << objectId;
            ctx.updateResult().nonExistentObjectIds.insert(objectId);
        }
    }
    return revIds;
}

std::string
categoryIds(const revision::Attributes& attributes)
{
    std::string result;
    for (const auto& pair : attributes) {
        if (pair.first.find(CATEGORY_PREFIX) == 0) {
            result += ((result.empty() ? "" : ",") + pair.first);
        }
    }
    return result;
}

void
updateAttributesBatch(
    GlobalContext& ctx,
    const std::vector<revision::DBID>& objectIds,
    const std::string& categoryId,
    const AttributesData& attributesData,
    const revision::Attributes& commitAttributes,
    revision::UserID updatedBy)
{
    auto revIds = loadRevisionIds(ctx, objectIds);
    if (revIds.empty()) {
        return;
    }

    const auto& gw = ctx.gateway();
    auto revs = gw.reader().loadRevisions(ctx.filterRevisionIds(std::move(revIds)));
    INFO() << revs.size() << " revisions loaded";

    std::list<revision::RevisionsGateway::NewRevisionData> newObjectsData;

    for (const auto& rev : revs) {
        const revision::ObjectRevision::Data& data = rev.data();
        const Attribute& newAttr = attributesData.at(rev.id().objectId());
        if (data.deleted) {
            INFO() << "Object already deleted " << rev.id();
            continue;
        }
        REQUIRE(data.attributes, "No attributes: " << rev.id());
        if (!data.attributes->count(categoryId)) {
            WARN() << "Not of " << categoryId << " category: " << rev.id()
                   << " " << categoryIds(*data.attributes);
            continue;
        }
        auto it = data.attributes->find(newAttr.name);
        if (it == data.attributes->end()) {
            INFO() << "Object " << rev.id()
                   << " set new value: " << newAttr.value;
        } else {
            if (it->second == newAttr.value) {
                INFO() << "Object already updated " << rev.id();
                continue;
            }
            if (!ctx.canOverwrite) {
                INFO() << "Object " << rev.id()
                       << " SKIP update new value: " << newAttr.value
                       << " was: " << it->second;
                continue;
            }
            INFO() << "Object " << rev.id()
                   << " set new value: " << newAttr.value
                   << " was: " << it->second;
        }
        if (ctx.dryRun) {
            continue;
        }
        revision::RevisionsGateway::NewRevisionData newObjData;
        newObjData.first = rev.id();
        newObjData.second.attributes = *data.attributes;
        (*newObjData.second.attributes)[newAttr.name] = newAttr.value;
        newObjData.second.deleted = false;
        newObjectsData.push_back(newObjData);
    }
    ctx.createCommit(newObjectsData, commitAttributes, updatedBy);
}

} // namespace

UpdateResult
updateRevisionAttributes(
    const std::string& connString,
    const std::string& categoryId,
    const AttributesData& attributesData,
    const revision::Attributes& commitAttributes,
    revision::UserID updatedBy,
    const Options& options)
{
    GlobalContext ctx(connString, options);
    auto batchSize = ctx.batchSize ? ctx.batchSize : BATCH_SIZE;

    std::list<revision::DBID> objectIds;
    std::transform(attributesData.begin(), attributesData.end(),
        std::inserter(objectIds, objectIds.end()),
        [] (const std::pair<revision::DBID, Attribute>& pair) -> revision::DBID
            { return pair.first; });

    std::vector<revision::DBID> batchObjectIds;

    auto updateBatch = [&]()
    {
        if (batchObjectIds.empty()) {
            return;
        }

        for (size_t pass = 0; ++pass <= ATTEMPTS; ) {
            try {
                updateAttributesBatch(
                    ctx, batchObjectIds, categoryId, attributesData,
                    commitAttributes, updatedBy);
                break;
            } catch (const std::exception& e) {
                ERROR() << "Attempt: " << pass << "/" << ATTEMPTS << " : " << e.what();
                if (pass == ATTEMPTS) {
                    throw;
                }
                ctx.refreshDraftCommits();
            }
        }
    };

    for (auto id : objectIds) {
        batchObjectIds.push_back(id);
        if (batchObjectIds.size() >= batchSize) {
            updateBatch();
            batchObjectIds.clear();
        }
    }
    updateBatch();

    auto& result = ctx.updateResult();
    INFO() << "Successfully updated " << result.updateCount << " objects";
    return std::move(result);
}
