#include <yandex/maps/wiki/pubsub/commit_consumer.h>

#include <yandex/maps/wiki/common/string_utils.h>

#include <boost/format.hpp>
#include <set>

namespace maps::wiki::pubsub {

namespace {

const size_t DEFAULT_BATCH_SIZE_LIMIT = 10;
const size_t DEFAULT_OUT_OF_ORDER_LIMIT = 10000;
const int64_t DEFAULT_TXID = 0;
const CommitId DEFAULT_COMMIT_ID = 0;

const std::string TXID = "txid";
const std::string COMMIT_ID = "commit_id";

const boost::format LOCK_WATERMARK_QUERY(
    "SELECT * FROM pubsub.lock_watermark(%1%, %2%)");

const boost::format SELECT_OUT_OF_ORDER_QUERY(
    R"(SELECT *
    FROM pubsub.out_of_order
    WHERE consumer_id=%1% AND branch_id=%2%)");

const std::string CURRENT_SNAPSHOT_TXMIN_QUERY =
    "SELECT txid_snapshot_xmin(txid_current_snapshot())";

const boost::format SELECT_IN_ORDER_QUERY(
    R"(SELECT txid, commit_id
    FROM revision.commit_queue
    WHERE branch_id=%1%
    AND (txid, commit_id) > (%2%))");

const boost::format SELECT_NEXT_BATCH_QUERY(
    R"(SELECT txid, commit_id
    FROM revision.commit_queue
    WHERE branch_id=%1%
    AND (txid, commit_id) > (%2%)
    ORDER BY txid, commit_id
    LIMIT %3%)");

const boost::format SELECT_WATERMARK_QUERY(
    R"(SELECT txid, commit_id
    FROM pubsub.watermark
    WHERE consumer_id=%1% AND branch_id=%2%)");

const boost::format UPDATE_WATERMARK_QUERY(
    R"(UPDATE pubsub.watermark
    SET txid=%1%, commit_id=%2%
    WHERE consumer_id=%3% AND branch_id=%4%)");

const boost::format DELETE_OLD_OUT_OF_ORDER_QUERY(
    R"(DELETE FROM pubsub.out_of_order
    WHERE consumer_id=%1% AND branch_id=%2% AND (txid, commit_id) <= (%3%))");

const boost::format INSERT_NEW_OUT_OF_ORDER_QUERY(
    R"(INSERT INTO pubsub.out_of_order
        (consumer_id, branch_id, txid, commit_id)
    VALUES(%1%))");

const boost::format SELECT_LAST_WATERMARK_OUT_OF_ORDER(
    R"(SELECT txid, commit_id
    FROM revision.commit_queue
    WHERE branch_id=%1%
    ORDER BY commit_id DESC
    LIMIT 1)");

const boost::format SELECT_LAST_WATERMARK(
    R"(SELECT txid, commit_id
    FROM revision.commit_queue
    WHERE txid <= (SELECT txid_snapshot_xmin(txid_current_snapshot())) AND branch_id=%1%
    ORDER BY commit_id DESC
    LIMIT 1)");

} // namespace

CommitConsumer::CommitConsumer(
        pqxx::transaction_base& consumerTxn,
        ConsumerId consumerId,
        BranchId branchId)
    : consumerTxn_(consumerTxn)
    , consumerId_(std::move(consumerId))
    , branchId_(branchId)
    , batchSizeLimit_(DEFAULT_BATCH_SIZE_LIMIT)
    , outOfOrderLimit_(DEFAULT_OUT_OF_ORDER_LIMIT)
{
    auto wmR = consumerTxn_.exec(
            (boost::format(LOCK_WATERMARK_QUERY)
             % consumerTxn_.quote(consumerId_) % branchId_).str());
    if (wmR.empty()) {
        throw AlreadyLockedException();
    }
    watermark_ = Item(wmR[0]);

    auto oooR = consumerTxn_.exec(
            (boost::format(SELECT_OUT_OF_ORDER_QUERY)
             % consumerTxn_.quote(consumerId_)
             % branchId_).str());
    for (const auto& row : oooR) {
        outOfOrderItems_.emplace(row);
    }
}

std::vector<CommitId>
CommitConsumer::consumeBatch(pqxx::transaction_base& revisionTxn)
{
    auto txmin = revisionTxn.exec(CURRENT_SNAPSHOT_TXMIN_QUERY)[0][0].as<int64_t>();

    auto offsetItem = watermark_;
    std::vector<Item> items;

    while (true) {
        auto result = revisionTxn.exec(
                (boost::format(SELECT_NEXT_BATCH_QUERY)
                % branchId_
                % offsetItem.toString()
                % batchSizeLimit_).str());

        for (const auto& row : result) {
            offsetItem = Item(row);

            if (outOfOrderItems_.count(offsetItem)) {
                continue; // skip
            }
            items.push_back(offsetItem);
        }

        // we have got all needed items
        if (items.size() == batchSizeLimit_) {
            break;
        }

        // no more commits in the queue
        if (result.size() < batchSizeLimit_) {
            break;
        }
    }

    std::vector<CommitId> batch;

    auto newWatermark = watermark_;
    auto itemIt = items.begin();
    for (; itemIt != items.end() && itemIt->txid <= txmin; ++itemIt) {
        newWatermark = *itemIt;
        batch.push_back(itemIt->key);
    }
    if (items.empty() || itemIt != items.end()) {
        // no more in-order items left.
        // we now can update new watermark using old out-of-order items.
        for (const auto& item : outOfOrderItems_) {
            if (item.txid <= txmin) {
                newWatermark = std::max(newWatermark, item);
            }
        }
    }
    updateWatermark(newWatermark);

    size_t ooCount = outOfOrderItems_.size();
    std::vector<Item> newOutOfOrder;
    for (; itemIt != items.end(); ++itemIt) {
        ooCount++;
        if (ooCount > outOfOrderLimit_) {
            break;
        }
        newOutOfOrder.push_back(*itemIt);
        batch.push_back(itemIt->key);
    }
    insertOutOfOrder(newOutOfOrder);

    return batch;
}

void CommitConsumer::updateWatermark(const Item& newWatermark)
{
    if (watermark_ == newWatermark) {
        return;
    }
    consumerTxn_.exec(
            (boost::format(UPDATE_WATERMARK_QUERY)
             % newWatermark.txid
             % newWatermark.key
             % consumerTxn_.quote(consumerId_)
             % branchId_).str());
    watermark_ = newWatermark;

    consumerTxn_.exec(
            (boost::format(DELETE_OLD_OUT_OF_ORDER_QUERY)
             % consumerTxn_.quote(consumerId_)
             % branchId_
             % newWatermark.toString()).str());
    outOfOrderItems_.erase(
            std::begin(outOfOrderItems_), outOfOrderItems_.upper_bound(newWatermark));
}

void CommitConsumer::insertOutOfOrder(
        const std::vector<Item>& newOutOfOrder)
{
    if (newOutOfOrder.empty()) {
        return;
    }

    auto ooStr = common::join(
            newOutOfOrder.begin(), newOutOfOrder.end(),
            [this](const Item& item) {
                return consumerTxn_.quote(consumerId_) + "," + std::to_string(branchId_)
                    + "," + item.toString();
            },
            "),(");
    consumerTxn_.exec(
            (boost::format(INSERT_NEW_OUT_OF_ORDER_QUERY) % ooStr).str());
    outOfOrderItems_.insert(newOutOfOrder.begin(), newOutOfOrder.end());
}

std::optional<CommitId> CommitConsumer::fastForwardToLastCommit()
{
    auto const& query = isOutOfOrderDisabled() ?
        SELECT_LAST_WATERMARK : SELECT_LAST_WATERMARK_OUT_OF_ORDER;
    auto wmR = consumerTxn_.exec((boost::format(query) % branchId_).str());
    if(wmR.empty()) {
        return std::nullopt;
    }

    updateWatermark(Item(wmR[0]));
    return wmR[0][COMMIT_ID].as<CommitId>();
}

bool CommitConsumer::isFirstRun() const {
    return watermark_.txid == DEFAULT_TXID
        && watermark_.key == DEFAULT_COMMIT_ID;
}

CommitConsumer& CommitConsumer::setBatchSizeLimit(size_t batchSizeLimit)
{
    REQUIRE(batchSizeLimit > 0, "batchSizeLimit should be > 0");
    batchSizeLimit_ = batchSizeLimit;
    return *this;
}

CommitConsumer& CommitConsumer::setOutOfOrderDisabled(bool value)
{
    outOfOrderLimit_ = (value ? 0 : DEFAULT_OUT_OF_ORDER_LIMIT);
    return *this;
}

size_t countCommitsInQueue(
    pqxx::transaction_base& consumerTxn,
    pqxx::transaction_base& revisionTxn,
    const ConsumerId& consumerId,
    BranchId branchId)
{
    std::set<Item> outOfOrderItems;
    auto oooResult = consumerTxn.exec(
            (boost::format(SELECT_OUT_OF_ORDER_QUERY)
             % consumerTxn.quote(consumerId)
             % branchId).str());
    for (const auto& row : oooResult) {
        outOfOrderItems.emplace(row);
    }

    auto watermarkResult = consumerTxn.exec(
        (boost::format(SELECT_WATERMARK_QUERY)
            % consumerTxn.quote(consumerId)
            % branchId).str());

    if (watermarkResult.empty()) {
        return 0;
    }

    Item batchStart{watermarkResult[0]};

    size_t count = 0;
    while (true) {
        static constexpr size_t BATCH_SIZE_LIMIT = 10000;
        auto itemsResult = revisionTxn.exec(
                (boost::format(SELECT_NEXT_BATCH_QUERY)
                 % branchId
                 % batchStart.toString()
                 % BATCH_SIZE_LIMIT).str());

        if (itemsResult.empty()) {
            break;
        }

        for (const auto& row : itemsResult) {
            batchStart = Item{row};
            if (outOfOrderItems.count(batchStart) == 0) {
                ++count;
            }
        }
    }

    return count;
}

} // namespace maps::wiki::pubsub
