#include "merger.h"

#include <dj/lib/processing/yt/proto/profile_yt.pb.h>
#include <dj/lib/profile/profile_map.h>
#include <dj/lib/shard2/index/shard_builder.h>

#include <robot/jupiter/library/rtdoc/file/docidmap_io.h>
#include <robot/jupiter/library/rtdoc/file/model/files.h>
#include <robot/jupiter/library/rtdoc/file/reader.h>
#include <robot/jupiter/library/rtdoc/file/writer.h>
#include <robot/jupiter/library/rtdoc/protos/lumps.pb.h>

#include <kernel/keyinv/indexfile/indexfile.h>
#include <kernel/keyinv/indexfile/indexwriter.h>
#include <kernel/tarc/docdescr/docdescr.h>
#include <kernel/tarc/iface/farcface.h>

#include <library/cpp/logger/global/global.h>

namespace NRTYServer {
    namespace {
        TString MainPrepPath(const TString& path) {
            return JoinFsPaths(path, NRtDoc::TBuilderFiles::MainPrepFile);
        }
    }

    TDjMerger::TDjMerger(const NDJ::NShard::TSearchIndexBuilder* builder)
        : Builder(builder)
    {
    }

    void TDjMerger::Merge(const TMergeContext& context) const {
        CHECK_WITH_LOG(context.Context.Decoder != nullptr);
        const TRTYMergerContextBase& mc = context.Context;
        const IRTYMergerDocIdDecoder& decoder = *mc.Decoder;
        if (mc.Sources.empty() || mc.Dests.empty() || mc.Decoder->GetNewDocsCount() == 0) {
            return;
        }

        Y_ENSURE(mc.Dests.size() == 1);
        const auto destCluster = 0;

        if (decoder.GetNewDocsCount(destCluster) == 0) {
            return;
        }

        const auto& destDir = mc.Dests.at(destCluster);

        NRtDoc::TPrepWriter prepWriter(MainPrepPath(destDir));
        TVector<ui16> indexFrq;

        for (ui32 sourceCl = 0; sourceCl < mc.Sources.size(); ++sourceCl) {
            NRtDoc::TPrepTable source(MainPrepPath(mc.Sources[sourceCl]));
            source.Init();
            for (auto it = source.Open(); it->Next(); ) {
                const auto [doc, tmpDocId, sourceDocId] = it->GetData();
                TRTYMergerAddress target = decoder.Decode(sourceCl, sourceDocId);
                if (target.IsRemoved()) {
                    continue;
                }
                Y_ENSURE(target.ClusterId == destCluster);
                const auto docId = target.DocId;

                prepWriter.AddData(docId, TBlob::NoCopy(doc.data(), doc.size()));

                if (docId >= indexFrq.size()) {
                    indexFrq.resize(docId + 1, -1);
                }

                indexFrq[docId] = 1;
            }
        }
        prepWriter.Finish();

        TFileOutput out(destDir + "/indexfrq");
        out.Write(reinterpret_cast<const char *>(indexFrq.data()), indexFrq.size() * sizeof(ui16));
        out.Finish();

        Builder->Build(destDir, destDir);
    }
}
