#pragma once

#include "assert_policy.h"
#include "bundle_meta.h"
#include "bundle_model.h"
#include "wad_patcher.h"
#include "log.h"
#include "timer_wrapper.h"

#include <robot/jupiter/library/rtdoc/protos/builder_task.pb.h>
#include <robot/jupiter/library/rtdoc/merger/counts_merger.h>
#include <robot/jupiter/library/rtdoc/merger/erf.h>
#include <robot/jupiter/library/rtdoc/file/portions_output.h>
#include <robot/jupiter/library/rtdoc/file/yt_client.h>
#include <saas/rtyserver_jupi/library/extbuilder/pudge.h>

#include <robot/jupiter/library/opt/mropt.h>
#include <robot/jupiter/library/opt/common/common.h>
#include <robot/jupiter/library/opt/shard_merge.h>
#include <mapreduce/yt/interface/client.h>
#include <util/generic/ptr.h>
#include <util/folder/path.h>

namespace NFusion {
    class TExtBuilderTask: public TThrRefBase {
    private:
        NRtDoc::TStopHandle::TPtr Stop;
        NRtDoc::TBuilderTask Task;
        NRtDoc::TBuilderTask DeltaTask;
        TString DeltaMapping;
        TMaybe<TBundleVersion> BundleVersion;
        THolder<TExtBuilderBundle> Bundle;
        NRtDoc::TBuilderPortionsConfig::TPtr Portions;
        bool WadDeltaMode;
        TPudgeData PudgeData;
        TVector<TStringBuf> PrepsToExclude;
        TString RebuildTag;
        bool IsRebuildMode = false;
        bool HeavyDataLumpsEnabled = false;

        NJupiter::TShardMergerOpts ShardsOpts;
        TAssertPolicy Asserts;
        TLog& Log;
        TTimerLog TimerLog;

        static const TString RebuildTagFileName;

    public:
        TExtBuilderTask(const NRtDoc::TBuilderTask& task, const NRtDoc::TStopHandle::TPtr& stop, const TAssertPolicy* assertPolicy, TLog* log)
            : Stop(stop)
            , Task(task)
            , Asserts(assertPolicy ? *assertPolicy : TAssertPolicy())
            , Log(log ? *log : TComponentLog::Default())
            , TimerLog(task.GetTimerLogPath(), task.GetRealm(), TFsPath(task.GetOutput().GetTrgDir()).GetName())
        {
            if (task.HasConfig() && task.GetConfig().GetDebugInfo() == true) {
                DebugDump(task);
            }

            Y_ENSURE_EX(task.HasOutput() && task.InputsSize() >= 1, yexception() << "incorrect TBuilderTask");
        }

    protected:
        virtual void Init(const NRtDoc::TBuilderConfig& cfg);

        virtual NRtDoc::TBuilderLocalClientPtr CreateLocalClient(const NRtDoc::IBuilderInputs& task, bool delta);

        virtual NRtDoc::TBuilderLocalClientPtr CreatePortionsClient(const NRtDoc::IBuilderInputs& task, bool delta);

        virtual NRtDoc::TBuilderLocalClientPtr CreateReducePortionsClient(const NRtDoc::IBuilderInputs& task, const TString& finalDir);

        virtual THolder<TExtBuilderBundle> GetBundle(const TBundleVersion& version, const TString& modelsDir, const TString& targetDir);

        virtual void InitDeltaTask();

        NRtDoc::TBuilderLocalClientPtr InitPreparePortions(bool delta, bool hostData = false);

        //
        // Migrations and legacy
        //
        void MigrateIndexes(TJupiterCmds& allCommands, const TSpecialCmds& specCommands);

        void WriteShardMetadata(const NRtDoc::TBuilderTask& builderTask);

        //
        // Multi-document preparates
        //
        void PreparePortions(const TMercuryCmds& allCommands, const TSpecialCmds& specCommands, bool delta);

        //
        // "Full rebuild" indexes
        //
        void MergePortions(const TMercuryCmds& allCommands, const TSpecialCmds& specCommands);

        void MergePortion(const TMercuryPrepCmd& command, NRtDoc::TSegmentHitRemapper::TPtr remapper);

        void ReduceMergedPortion(const TMercuryPrepCmd& command);

        void MergePreparates();

        void BuildFullIndexes(const TJupiterCmds& allCommands, const TSpecialCmds& specCommands);

        //
        // "Wad merge" indexes
        //
        void BuildWadDeltas(const TJupiterCmds& allCommands);

        NRtDoc::IWadPatcher::TPtr CreateWadPatcher(const TJupiterCmd& command, const TString& indexaaPrefix);

        static NRtDoc::TBuilderTask MakeMergerTask(const NRtDoc::TBuilderTask& builderTask, const NRtDoc::TBuilderTask::TBuilderInput& deltaInput);

        NRtDoc::IWadMerger::TPtr CreateWadMerger(const TJupiterCmd& command);

        void MergeWads(const TJupiterCmds& allCommands);

        void BuildAndMergeWads(const TJupiterCmds& allCommands);

        //
        // "Chunked wad" indexes
        //
        void PudgeWads(const TJupiterCmds& allCommands);

    private:
        // Debug routines
        void DebugDump(const NRtDoc::TBuilderTask& proto);

        void RunShardsPrepare();
        void GenerateNewPudgeMapping();
        void RemovePreviouslyUpdatedPrepLumps();

        void InitRebuildMode();
        TString LoadRebuildTag() const;
        TString ReadSegmentRebuildTag(const TString& dir) const;
        void WriteRebuildTagIfNeed() const;

    public:
        //sync run
        bool Run(NRtDoc::TBuilderTaskResult& /*result*/);
    };
}
