#include "utils.h"

#include <maps/libs/chrono/include/days.h>
#include <maps/libs/cmdline/include/cmdline.h>
#include <maps/libs/common/include/exception.h>
#include <maps/libs/log8/include/log8.h>
#include <maps/libs/pgpool/include/pgpool3.h>
#include <maps/wikimap/mapspro/libs/acl/include/aclgateway.h>
#include <maps/wikimap/mapspro/libs/acl/include/user.h>
#include <maps/wikimap/mapspro/libs/gdpr/include/user.h>
#include <yandex/maps/wiki/common/default_config.h>
#include <yandex/maps/wiki/common/extended_xml_doc.h>
#include <yandex/maps/wiki/common/moderation.h>
#include <yandex/maps/wiki/common/pgpool3_helpers.h>
#include <yandex/maps/wiki/common/string_utils.h>
#include <yandex/maps/wiki/common/pg_advisory_lock_ids.h>
#include <yandex/maps/wiki/configs/editor/category_groups.h>
#include <yandex/maps/wiki/configs/editor/config_holder.h>
#include <yandex/maps/wiki/pubsub/commit_consumer.h>
#include <yandex/maps/wiki/revision/commit.h>
#include <yandex/maps/wiki/revision/filters.h>
#include <yandex/maps/wiki/social/gateway.h>
#include <yandex/maps/wiki/tasks/status_writer.h>


#include <yandex/maps/pgpool3utils/pg_advisory_mutex.h>

#include <pqxx/pqxx>

#include <algorithm>
#include <chrono>
#include <iostream>
#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

namespace acl = maps::wiki::acl;
namespace common = maps::wiki::common;
namespace revision = maps::wiki::revision;
namespace rf = maps::wiki::revision::filters;
namespace cfgeditor = maps::wiki::configs::editor;

namespace {

const std::string TASK_TYPE = "update-rating";

const std::string CORE_DB_ID = "core";
const std::string SOCIAL_DB_ID = "social";
const std::string POOL_ID = "grinder";

const std::string RATING_TYPE_ARG = "rating-type";

using Clock = std::chrono::system_clock;

const size_t INITIAL_BATCH_SIZE = 1000;
const std::string CONSUMER_ID_PREFIX = "DelayedStats::";

const size_t WRITE_BULK_SIZE = 100;

const std::string COUNTERS_V2_SUFFIX = "_v2";

struct RatingTraits
{
    std::string id;
    std::string delayedCounterName;
    std::chrono::seconds counterDelay;
};

const std::string CURRENT_EDITS_COUNTER = "total_edits";

} // namespace

namespace maps::wiki::rating {

const std::unordered_map<std::string, RatingTraits>
RATING_TRAITS = {
    {"full", {"full", {}, {}}},
    {"day", {"day", "edits_1d_ago", chrono::Days{1}}},
    {"week", {"week", "edits_7d_ago", chrono::Days{7}}},
    {"month", {"month", "edits_30d_ago", chrono::Days{30}}},
    {"quarter", {"quarter", "edits_90d_ago", chrono::Days{90}}}
};

using UserToCount = std::map<revision::UserID, size_t>;

UserToCount
commitsCountByAuthor(
    const social::Events& commitEvents,
    const cfgeditor::CategoryGroups& categoryGroups)
{
    auto isEventToCount = [&](const social::Event& event) {
        if (auto category = event.getPrimaryObjectCategory(); category) {
            auto group = categoryGroups.findGroupByCategoryId(*category);
            return group && group->id() != "service_group";
        }
        return event.action().starts_with("group");
    };

    UserToCount userToCount;
    for (const auto& event : commitEvents) {
        const gdpr::User user(event.createdBy());
        if (!user.hidden() && isEventToCount(event)) {
            userToCount[user.uid()]++;
        }
    }
    return userToCount;
}

social::Events loadCommitEvents(
    pqxx::transaction_base& socialTxn,
    const std::vector<revision::DBID>& commitIds,
    chrono::TimePoint createdBefore)
{
    social::Gateway gtw(socialTxn);
    social::TIds commitIdsSet(commitIds.begin(), commitIds.end());
    return gtw.loadTrunkEditEventsByCommitIdsCreatedBefore(commitIdsSet, createdBefore);
}

class StatsCounter
{
public:
    StatsCounter(pqxx::transaction_base& txn, std::string columnName)
        : txn_(txn)
        , columnName_(std::move(columnName))
    {}

    bool increase(revision::UserID uid, size_t amount)
    {
        std::ostringstream query;
        query <<
            "INSERT INTO social.stats AS stats (uid, " << columnName_ << ") "
            "VALUES (" << uid << ", " << amount << ") "
            "ON CONFLICT (uid) DO "
            "UPDATE SET " <<
                columnName_ << " = COALESCE(stats." << columnName_ << ", 0) + " << amount << " "
                "WHERE stats.uid = " << uid;

         return txn_.exec(query.str()).affected_rows();
    }

private:
    pqxx::transaction_base& txn_;
    std::string columnName_;
};

void updateUsersStatsCounters(
    pqxx::transaction_base& socialTxn,
    const std::string& counterName,
    const std::map<revision::UserID, size_t>& countsByAuthor)
{
    StatsCounter counter(socialTxn, counterName + COUNTERS_V2_SUFFIX);
    for (const auto& [uid, count] : countsByAuthor) {
        if (!counter.increase(uid, count)) {
            throw maps::Exception()
                << "cannot update delayed counters for user " << uid
                << ", stats row does not exist";
        }
    }
}

void updateDelayedCounters(
    maps::pgpool3::Pool& corePool,
    maps::pgpool3::Pool& socialPool,
    const cfgeditor::CategoryGroups& categoryGroups,
    const RatingTraits& ratingTraits)
{
    if (ratingTraits.delayedCounterName.empty()) {
        return;
    }

    INFO() << "Updating " << ratingTraits.delayedCounterName << " counters";

    const auto createdAtBefore = Clock::now() - ratingTraits.counterDelay;
    const auto consumerId = CONSUMER_ID_PREFIX + ratingTraits.delayedCounterName;

    auto batchSizeLimit = INITIAL_BATCH_SIZE;
    size_t totalCommitsProcessed = 0;
    while (batchSizeLimit) {
        auto socialTxn = socialPool.masterWriteableTransaction();
        maps::wiki::pubsub::CommitConsumer consumer(
            *socialTxn, consumerId, revision::TRUNK_BRANCH_ID);
        consumer.setBatchSizeLimit(batchSizeLimit);

        auto coreTxn = corePool.slaveTransaction();
        auto batch = consumer.consumeBatch(*coreTxn);
        if (batch.empty()) {
            break;
        }

        const auto commitIds = revision::Commit::loadIds(*coreTxn,
            rf::CommitAttr::id().in(batch) &&
            rf::CommitCreationTime() < Clock::to_time_t(createdAtBefore));

        if (commitIds.size() < batch.size()) {
            batchSizeLimit = commitIds.size();
            continue;
        }

        auto commitEvents = loadCommitEvents(*socialTxn, batch, createdAtBefore);

        updateUsersStatsCounters(
            *socialTxn,
            ratingTraits.delayedCounterName,
            commitsCountByAuthor(commitEvents, categoryGroups)
        );

        totalCommitsProcessed += commitIds.size();
        DEBUG() << "Processed " << totalCommitsProcessed << " new commits so far";
        socialTxn->commit();
    }
    INFO() << "Processed " << totalCommitsProcessed << " new commits";
}

struct UserRating
{
    UserRating(acl::UID uid_, size_t editsCount_)
        : uid(uid_)
        , editsCount(editsCount_)
    { }

    acl::UID uid;
    size_t editsCount;
};

std::vector<UserRating> calculateRating(
    pqxx::transaction_base& txn,
    const RatingTraits& ratingTraits,
    const std::unordered_set<acl::UID>& skipUids)
{
    auto editsCounter = CURRENT_EDITS_COUNTER + COUNTERS_V2_SUFFIX;
    if (ratingTraits.id != "full") {
        auto columnName = ratingTraits.delayedCounterName + COUNTERS_V2_SUFFIX;
        editsCounter += " - COALESCE(" + columnName + ", 0)";
    }

    auto query =
        "SELECT uid, " + editsCounter +
        " FROM social.stats WHERE " + editsCounter + " > 0";

    std::vector<UserRating> rating;
    for (const auto& row : txn.exec(query)) {
        auto uid = row[0].as<acl::UID>();
        if (!skipUids.count(uid)) {
            rating.emplace_back(uid, row[1].as<size_t>());
        }
    }

    std::sort(
        std::begin(rating), std::end(rating),
        [](const UserRating& lhs, const UserRating& rhs)
        {
            if (lhs.editsCount == rhs.editsCount) {
                return lhs.uid < rhs.uid;
            }
            return lhs.editsCount > rhs.editsCount;
        });

    return rating;
}

void writeRating(
    pqxx::transaction_base& txn,
    const RatingTraits& ratingTraits,
    const std::vector<UserRating>& rating)
{
    auto tableNameShort = ratingTraits.id + COUNTERS_V2_SUFFIX;
    auto tableName = "rating." + tableNameShort;
    auto tableNameTmp = tableName + "_tmp";

    txn.exec(
        "CREATE TABLE " + tableNameTmp +
        " (LIKE " + tableName + " INCLUDING ALL)");

    size_t pos = 1;
    auto ratingFormatter = [&pos](const UserRating& r) {
        std::ostringstream os;
        os << '(' << pos << ',' << r.uid << ',' << r.editsCount << ')';
        ++pos;
        return os.str();
    };

    auto bulkBegin = std::begin(rating);
    while (bulkBegin != std::end(rating)) {
        auto bulkEnd = std::next(
            bulkBegin,
            std::min<size_t>(
                WRITE_BULK_SIZE,
                std::distance(bulkBegin, std::end(rating))));

        txn.exec(
            "INSERT INTO " + tableNameTmp +
            " (pos, uid, score) VALUES " +
            common::join(bulkBegin, bulkEnd, ratingFormatter, ','));

        bulkBegin = bulkEnd;
    }

    txn.exec(
        "DROP TABLE " + tableName + ";"
        "ALTER TABLE " + tableNameTmp +" RENAME TO " + tableNameShort);
}

void updateRatingMeta(
    pqxx::transaction_base& txn,
    const RatingTraits& ratingTraits,
    const std::vector<UserRating>& rating,
    time_t now)
{
    auto ratingMetaTable = "rating.meta" + COUNTERS_V2_SUFFIX;

    txn.exec("UPDATE " + ratingMetaTable + " "
        "SET version = " + std::to_string(now) +
            ", size = " + std::to_string(rating.size()) +
        " WHERE type = " + txn.quote(ratingTraits.id));
}

void updateRating(
        maps::pgpool3::Pool& socialPool,
        const RatingTraits& ratingTraits,
        const std::unordered_set<acl::UID>& skipUids)
{
    INFO() << "Updating " << ratingTraits.id << " rating";

    auto now = Clock::to_time_t(Clock::now());

    auto txn = socialPool.masterWriteableTransaction();
    auto rating = calculateRating(*txn, ratingTraits, skipUids);

    writeRating(*txn, ratingTraits, rating);
    INFO() << "Wrote rating " << ratingTraits.id
           << ", version " << now << ", size " << rating.size();

    updateRatingMeta(*txn, ratingTraits, rating, now);
    INFO() << "Updated metadata for " << ratingTraits.id << " rating";

    txn->commit();
}

bool run(const common::ExtendedXmlDoc& configDoc)
{
    common::PoolHolder socialPoolHolder(configDoc, SOCIAL_DB_ID, POOL_ID);

    pgp3utils::PgAdvisoryXactMutex locker(
        socialPoolHolder.pool(),
        static_cast<int64_t>(common::AdvisoryLockIds::RATING_WORKER));
    if (!locker.try_lock()) {
        INFO() << "Database is already locked. Task interrupted.";
        return false;
    }

    common::PoolHolder corePoolHolder(configDoc, CORE_DB_ID, POOL_ID);

    cfgeditor::ConfigHolder editorConfig(
        configDoc.get<std::string>("/config/services/editor/config")
    );

    auto skipUids = fetchSkipUids(corePoolHolder.pool());

    for (const auto& [_, ratingTraits] : RATING_TRAITS) {
        updateDelayedCounters(
            corePoolHolder.pool(),
            socialPoolHolder.pool(),
            editorConfig.categoryGroups(),
            ratingTraits);

        updateRating(
            socialPoolHolder.pool(),
            ratingTraits,
            skipUids);
    }

    return true;
}

} // namespace maps::wiki::rating

int main(int argc, char** argv)
{
    try {
        maps::cmdline::Parser parser;

        auto syslogTag = parser.string("syslog-tag").help(
            "redirect log output to syslog with given tag");
        auto workerConfig = parser.file("config").help(
            "path to services configuration");
        auto statusDir = parser.string("status-dir").help(
            "path to status dir");

        parser.parse(argc, argv);

        if (syslogTag.defined()) {
            maps::log8::setBackend(maps::log8::toSyslog(syslogTag));
        }

        std::unique_ptr<common::ExtendedXmlDoc> configDocPtr;
        if (workerConfig.defined()) {
            configDocPtr.reset(new common::ExtendedXmlDoc(workerConfig));
        } else {
            configDocPtr = common::loadDefaultConfig();
        }

        auto success = maps::wiki::rating::run(*configDocPtr);

        if (success && statusDir.defined()) {
            maps::wiki::tasks::StatusWriter statusWriter(statusDir + "/wiki-rating-worker.status");
            statusWriter.flush();
        }

        return success
            ? EXIT_SUCCESS
            : EXIT_FAILURE;
    } catch (const maps::Exception& ex) {
        FATAL() << "Worker failed: " << ex;
        return EXIT_FAILURE;
    } catch (const std::exception& ex) {
        FATAL() << "Worker failed: " << ex.what();
        return EXIT_FAILURE;
    }
}
