#include "frq_merge.h"

#include <robot/jupiter/library/opt/inproc_args.h>
#include <robot/jupiter/library/tables/shards_prepare.h>
#include <robot/jupiter/library/utils/time_gatherer.h>
#include <robot/library/yt/static/table.h>
#include <robot/library/utils/commons.h>
#include <robot/library/utils/progress_logger.h>

namespace NFusion {

    int MergeFrq(const NJupiter::TInProcArgs& args) {
        const NJupiter::TShardMergerOpts& shardsOpts = args.ShardMergerOpts;
        NYT::IIOClient& client = args.Client;
        NJupiter::TShardsPrepareTables tables = NJupiter::TShardsPrepareTables::Clientless();
        NJupiter::TTimeGatherer timeGatherer(shardsOpts.StatsDir, __FUNCTION__);

        const TString outputPath = TFsPath(shardsOpts.OutputDir) / shardsOpts.FilePrefix;
        TFileOutput out(outputPath);
        out.SetFinishPropagateMode(true);

        NYT::TRichYPath tablePath = tables.GetInvUrlHashTable(/*bucketNum*/0).GetPath();
        NYT::TTableReaderPtr<NJupiter::TInvUrlHash> reader =
            NRobot::CreateSelectedShardTableReader<NJupiter::TInvUrlHash>(client, tablePath, shardsOpts.Shard);

        INFO_LOG << "Start writing " << shardsOpts.FilePrefix << " for " << shardsOpts.Shard << Endl;

        NJupiter::TProgressLogger prLog("docs");

        ui32 currentDocId = 0;
        const ui16 presentDoc = 1;
        const ui16 removedDoc = -1;
        for (; reader->IsValid(); reader->Next()) {
            const auto& row = reader->GetRow();
            const ui32 docId = row.GetDocId();
            for (; currentDocId < docId; ++currentDocId) {
                prLog.Step();
                out.Write(&removedDoc, sizeof(removedDoc));
            }
            Y_ENSURE(currentDocId == docId);
            prLog.Step();
            out.Write(&presentDoc, sizeof(presentDoc));
            ++currentDocId;
        }

        NJupiter::TTable<NJupiter::TUrlToShardStats> statTable = tables.GetStatsTableLink();
        NYT::TRichYPath shardPath = statTable.GetPath();
        NYT::TTableReaderPtr<NJupiter::TUrlToShardStats> curPtr =
            NRobot::CreateSelectedShardTableReader<NJupiter::TUrlToShardStats>(client, shardPath, shardsOpts.Shard);
        Y_ENSURE(curPtr->IsValid() && curPtr->GetRow().GetShard() == shardsOpts.Shard);
        const ui32 docCount = curPtr->GetRow().GetDocCount();
        for (; currentDocId < docCount; ++currentDocId) {
            prLog.Step();
            out.Write(&removedDoc, sizeof(removedDoc));
        }
        Y_ENSURE(currentDocId == docCount);

        out.Finish();
        prLog.Finish();
        INFO_LOG << "Finished writing " << shardsOpts.FilePrefix << " for " << shardsOpts.Shard << Endl;
        return 0;
    }

}
