#include "mingeo_component.h"
#include "mingeo_index.h"

#include <saas/rtyserver/common/should_stop.h>
#include <saas/rtyserver/components/l2/l2_memory_manager.h>

using namespace NRTYServer;

TMinGeoComponent::TMinGeoComponent(const TRTYServerConfig& config)
    : TBase(config)
{
    IndexFiles.clear();
    const TString invFileName = "index" + TString(NRTYGeo::TGeoIndexFormatter::Suffix); // "indexgeo.rty"
    IndexFiles.insert(TIndexFile(invFileName, /*check hash=*/true, TIndexFile::ppDisable));
}

bool TMinGeoComponent::DoMerge(const TMergeContext& context) const {
    if (!TBase::DoMerge(context))
        return false;

    using namespace NRTYGeo;

    struct TDest {
        TGeoIndexBuilder Builder;
        TFsPath File;
    };

    TVector<TFsPath> sourceInvs;
    for (ui32 src = 0; src < context.Context.Sources.size(); ++src) {
        TString srcDir = context.Context.Sources[src];
        TFsPath srcPath = TGeoIndexFormatter::FormatFileName(TFsPath(srcDir), "index");

        if (!srcPath.Exists() || !srcPath.IsFile()) {
            // This never happens, because the Normalizer should have taken care of this.
            ERROR_LOG << GetName() << ": file " << srcPath << " does not exist, skipping merge" << Endl;
            // No Geo functionality will be available
            return true;
        }

        sourceInvs.push_back(srcPath);
    }

    TVector<THolder<TDest>> builders;
    for (ui32 dst = 0; dst < context.Context.Dests.size(); ++dst) {
        THolder<TDest> dest = MakeHolder<TDest>();
        dest->File = TGeoIndexFormatter::FormatFileName(TFsPath(context.Context.Dests[dst]), "index");

        builders.emplace_back(std::move(dest));
    }

    for (ui32 src = 0; src < context.Context.Sources.size() && !ShouldStop(context.RigidStopSignal); ++src) {
        TFsPath srcPath = sourceInvs[src];
        TGeoIndex srcIndex;
        TGeoIndexFormatter::Load(srcPath, srcIndex);
        for (auto&& [partKey, pSrcTree] : srcIndex.Trees) {
            const auto& srcHits = pSrcTree->GetData();
            for (const TGeoIndexTree::TRecord& record : srcHits) {
                const ui32 docId = record.DocId;
                const auto& item = record.Coords;
                Y_ENSURE(TGeoIndexTreeBuilder::IsItem(item), "Incorrect data for docId " << docId << " in " << srcPath);

                TRTYMerger::TAddress addr = context.Context.Decoder->Decode(src, docId);
                if (Y_LIKELY(addr.DocId != REMAP_NOWHERE)) {
                    builders[addr.ClusterId]->Builder.AddHit(partKey.first, partKey.second, addr.DocId, item);
                }
            }
        }
    }

    for (ui32 dst = 0; dst < context.Context.Dests.size(); ++dst) {
        TGeoIndex dstIndex;
        builders[dst]->Builder.Finalize(dstIndex);
        TGeoIndexFormatter::Save(builders[dst]->File, dstIndex);
    }

    return true;
}
