#pragma once

#include <saas/rtyserver/config/config.h>
#include <saas/rtyserver/config/merger_config.h>
#include <saas/rtyserver/config/realm_config.h>
#include <saas/rtyserver/common/common_messages.h>
#include <saas/rtyserver/common/message_collect_server_info.h>
#include <saas/rtyserver/index_storage/index_storage.h>
#include <saas/rtyserver/merger/library/merger.h>
#include "aggregator.h"

class TMergeTask: public IMergerTask, public TRTYMerger::IRTYMergerDocIdInfo, public IMessageProcessor {
private:
    const int ShardNumber;
    const TVector<TString> Segments;
    const TRTYServerConfig& Config;
    IMergerTaskNotifier& Notifier;
    TMergerMetrics& Metrics;
    const bool IsPortions;

    THolder<TRTYMerger::IRTYMergerDocIdDecoder> Decoder;
    TVector<TString> DestDirs;
    TIndexStorage* Storage;
    THolder<TDeferredUpdatesStorage> UpdatesStorage;
    TString IndexMergeDir;
    TVector<THolder<NRTYMerger::TAggregator>> Aggregators;

protected:
    TString DoBuildTempDest(IIndexStorage& storage, const TString& segment) const override {
        return storage.GetFullPath(segment + "_merge", GetPath()) + GetDirectorySeparator();
    }

    TString DoBuildTmpfsDest(IIndexStorage& storage, const TString& segment) const override {
        if (Config.GetMergerConfig().TmpfsDir && GetSourceSegments().size() <= Config.GetMergerConfig().MaxPrepsToUseTmpfs) {
            return storage.GetFullPath(segment + "_merge_tmpfs", Config.GetMergerConfig().TmpfsDir) + GetDirectorySeparator();
        } else {
            return IMergerTask::DoBuildTmpfsDest(storage, segment);
        }
    }

    TString DoBuildFullDest(IIndexStorage& storage, const TString& segment) const override {
        return storage.GetFullPath(segment, GetPath()) + GetDirectorySeparator();
    }

    void DoBuildDecoder(IIndexStorage& storage) override;

    NJson::TJsonValue GetJsonInfo() const {
        NJson::TJsonValue result(NJson::JSON_MAP);
        NJson::TJsonValue sourcesInfo;
        NJson::TJsonValue destsInfo;
        for (ui32 i = 0; i < Segments.size(); ++i) {
            sourcesInfo.AppendValue(Segments[i]);
        }
        for (ui32 i = 0; i < DestDirs.size(); ++i) {
            destsInfo.AppendValue(DestDirs[i]);
        }
        result.InsertValue("merger_dir", IndexMergeDir);
        result.InsertValue("sources", sourcesInfo);
        result.InsertValue("destinations", destsInfo);
        AddJsonInfo(result);
        return result;
    }

    void ReportStart();
    void ReportEnd(bool success);

public:
    bool Process(IMessage* message) override;

    TString Name() const override {
        return "MergerTask";
    }

    bool DoOnStart(const std::atomic<bool>* rigidStop) override;
    void DoOnBeforeMainStage() override;
    void DoOnFailed() override;
    void DoOnFinished(const std::atomic<bool>* rigidStop) override;

    void OnBeforeUpdatableMerging() override {
        if (!UpdatesStorage) {
            CHECK_WITH_LOG(Storage);
            UpdatesStorage.Reset(new TDeferredUpdatesStorage(Config));
            for (TVector<TString>::const_iterator i = GetSourceSegments().begin(); i != GetSourceSegments().end(); ++i)
                Storage->GetFinalIndex(*i)->SetDeferredUpdatesStorage(UpdatesStorage.Get());
            SendGlobalDebugMessage<TMessageDeferredUpdaterActivated>(TMessageDeferredUpdaterActivated::EDeferredUpdateMoment::dumMerge);
        }
    }

    void MoveFromTemp(ui32 destIndex, IIndexStorage& storage, const std::atomic<bool>* rigidStop) override;

    ui64 Timestamp(ui32 clusterId, ui32 docId) const override;

    TRTYMerger::IRTYMergerDocIdDecoder* GetDecoder() override {
        return Decoder.Get();
    }

    TRTYMerger::IRTYMergerDocIdInfo* GetInfo() override {
        return this;
    }

    TMergeTask(int shardNumber, const TVector<TString>& segments, const TRTYServerConfig& config, const TString& path, IMergerTaskNotifier& notifier, TMergerMetrics& metrics, IMergerLockPolicy::TPtr lockPolicy, bool isPortions = false)
        : IMergerTask(path, lockPolicy, config.GetRealmListConfig().GetRealmConfig(path).ConfigName)
        , ShardNumber(shardNumber)
        , Segments(segments)
        , Config(config)
        , Notifier(notifier)
        , Metrics(metrics)
        , IsPortions(isPortions)
        , Storage(nullptr)
    {
        RegisterGlobalMessageProcessor(this);
        Notifier.OnMergerTaskStart(this);
    }

    ~TMergeTask() {
        UnregisterGlobalMessageProcessor(this);
        INFO_LOG << "Finished merger task info: " << GetJsonInfo().GetStringRobust() << Endl;
        Notifier.OnMergerTaskFinish(this);
    }

    int GetShardNumber() const override { return ShardNumber; }

    const TVector<TString>& GetSourceSegments() const override { return Segments; }
    const TVector<TString>& GetDestSegments() const override;
    TString GetName() const override {
        return "Merge task: " + JoinStrings(GetSourceSegments(), ",") + "->" + JoinStrings(GetDestSegments(), ",");
    };

    bool NeedWriteSourceIndexes() const override {
        return true;
    }

    bool GetIsPortions() const override {
        return IsPortions;
    }
};
