#pragma once

#include <robot/jupiter/library/opt/shard_merge.h>
#include <robot/jupiter/library/pudge/config/load.h>
#include <robot/jupiter/library/pudge/config/tier_config.h>
#include <robot/jupiter/library/rtdoc/file/yt_client.h>
#include <robot/jupiter/library/rtdoc/protos/builder_task.pb.h>
#include <robot/jupiter/library/pudge/minipudge/protos/pudge_chunk_stats.pb.h>
#include <robot/jupiter/protos/pudge_delta_chunk.pb.h>
#include <saas/rtyserver_jupi/library/extbuilder/pudge_mapping.pb.h>

#include <kernel/doom/chunked_wad/doc_chunk_mapping.h>

#include <util/generic/maybe.h>
#include <util/generic/vector.h>

namespace NFusion {

    class TPudgeData {
    public:
        using TCreateClient = std::function<NRtDoc::TBuilderLocalClientPtr(const TString&)>;

        void GenerateNewMapping(const NRtDoc::TBuilderTask& task, const TCreateClient& createClient);

        TString GetShardsOptsPrevIndexDir() const {
            return PrevFinal ? PrevFinal->GetSrcDir() : NJupiter::TShardMergerOpts::EmptyDir;
        }

        static TVector<NDoom::TDocChunkMapping> LoadDocChunkMapping(const TString& finalDir);

        static const NJupiter::TTierConfig LoadTierConfig();

    private:
        struct TDoc {
            ui32 DocId = 0;
            TMaybe<ui32> PrevDocId;
            TMaybe<ui32> Chunk;
            TMaybe<ui32> LocalId;
        };

        void GeneratePrevFinalPudgeMapping(NYT::IIOClient& client, const NRtDoc::TBuilderTask::TBuilderInput& prevFinal);

        TVector<ui32> SelectDeltaChunks(const TVector<TDoc>& docs) const;

        TVector<NJupiter::TPudgeDeltaChunk> AssignDocIds(
            NYT::IIOClient& client,
            const TVector<ui32>& pudgeDeltaChunkIndices,
            TVector<TDoc>& docs
        );

        void WritePudgeMapping(const TVector<TDoc>& docs, const TString& outputDir) const;

        static TVector<TDoc> CreateDocs(size_t count);

        static void SetChunkInfoFromPrevMapping(const NRtDoc::TBuilderTask::TBuilderInput& prevFinal, const TPudgeMapping& prevMapping, TVector<TDoc>& docs);

        TVector<NJupiter::TPudgeChunkStats> FillChunkStats(const TVector<TDoc>& docs) const;
        void ExtractWalrusArc(NYT::IIOClient& client, const TVector<TDoc>& docs) const;

        void WriteDeltaChunkTable(NYT::IIOClient& client, const TVector<NJupiter::TPudgeDeltaChunk>& deltaChunks) const;
        void WriteDocChunkMappingTables(NYT::IIOClient& client, const TVector<TDoc>& docs) const;

    public:
        static const TString ShardName;

    private:
        NJupiter::TTierConfig Config;
        TMaybe<NRtDoc::TBuilderTask::TBuilderInput> PrevFinal;

        static const TString MappingFileName;
    };

}
