#pragma once

#include <saas/searchproxy/proxy_meta/rearrange/abstract/rearrange.h>
#include <saas/library/hash_to_block_mode/hash_to_block_mode.h>
#include <saas/protos/rtyserver.pb.h>

#include <library/cpp/mediator/global_notifications/system_status.h>
#include <library/cpp/cgiparam/cgiparam.h>

#include <search/idl/meta.pb.h>

class TBlockModeRearrange : public ICustomRearrange {
    ui64 BlockCount;

public:
    TBlockModeRearrange(const ui64 blockCount) : ICustomRearrange(), BlockCount(blockCount) {
    }

    TString GetName() const override {
        return "BlockMode";
    }

    bool DoFormCgiParameters(TCgiParameters& cgi) override;

    void DoProcessReplies(TVector<NProxyMeta::TSourceReply>& replies, ICustomReportBuilder& /*builder*/, IReplyContext& /*context*/) const override;

    void FillDocument(NMetaProtocol::TDocument* newDoc, NRTYServer::TMessage::TDocument& doc) const;
};

class TBlockModeRearrangeFactory : public ICustomRearrangeFactory {
    ui64 BlockCount;

    static ICustomRearrangeFactory::TFactory::TRegistrator<TBlockModeRearrangeFactory> Registrator;
public:
    TBlockModeRearrangeFactory(const TCustomRearrangeParams& params) {
        ui64 TotalDocCount = 1, MaxBlockDocCount = 1;

        auto it = params.Parameters.find("TotalDocCount");
        if (it != params.Parameters.end()) {
            TotalDocCount = FromString<ui64>(it->second);
        }
        else {
            AbortFromCorruptedConfig("Incorrect TotalDocCount");
        }

        it = params.Parameters.find("MaxBlockDocCount");
        if (it != params.Parameters.end()) {
            MaxBlockDocCount = FromString<ui64>(it->second);
        }
        else {
            AbortFromCorruptedConfig("Incorrect MaxBlockDocCount");
        }

        BlockCount = (TotalDocCount + MaxBlockDocCount - 1) / MaxBlockDocCount;
    }

    TString GetName() const override {
        return "BlockMode";
    }

    virtual ICustomRearrange::TPtr CreateRerrangeInstance() const override {
        return new TBlockModeRearrange(BlockCount);
    }
};

ICustomRearrangeFactory::TFactory::TRegistrator<TBlockModeRearrangeFactory> TBlockModeRearrangeFactory::Registrator("BlockMode");
