#pragma once

#include "mingeo_core.h"
#include "mingeo_index.h"
#include "geo_kdtree.h"

#include <saas/rtyserver/components/l2/l2_builder.h>
#include <util/system/mutex.h>

namespace NRTYServer {
    class TMinGeoBuilder: public TL2ComponentBuilder {
        using TL2ComponentBuilder::TL2ComponentBuilder;

    private:
        TMutex Mutex;
        NRTYGeo::TGeoIndexBuilder Builder;

    public:
        Y_FORCE_INLINE static void Index(NRTYGeo::TGeoIndexBuilder& bld, const TGeoTransientEntity::TBuilderInvData& invRecord, const ui32 docId) {
            for (const auto& key : invRecord.Keys) {
                // AddDoc will add one or more (key,hit) pairs that represent (docId, TGeoWindow)
                bld.AddDoc(/*streamId=*/0, /*kps=*/invRecord.KeyPrefix, docId, NGeo::TGeoWindow(key.RectLower, key.RectUpper));
                //TODO(SAAS-5949): support multiple streamId per invRecord
            }
        }

        Y_FORCE_INLINE static void Index(NRTYGeo::TGeoIndexBuilder& bld, const TGeoTransientEntity* entity, const ui32 docId) {
            // the entity that we are going to index should be either obtained from Parse, or prepared via RestoreForMerger.
            if (!entity || !entity->BuilderExtension.Defined())
                return;

            const TGeoTransientEntity::TBuilderInvData& rect = *entity->BuilderExtension;
            Index(bld, rect, docId);
        }

        virtual void Index(int /*threadId*/, const TParsedDocument& doc, const ui32 docId) override {
            TGuard<TMutex> g(Mutex);
            auto* entity = doc.GetComponentEntity<TGeoTransientEntity>(ComponentName);
            Index(Builder, entity, docId);
        }

        virtual bool DoClose(const TBuilderCloseContext& context) override {
            TGuard<TMutex> g(Mutex);
            if (context.RemapTable) {
                Builder.Remap(*context.RemapTable);
            }
            NRTYGeo::TGeoIndex index;
            Builder.Finalize(index);

            const TFsPath outputPath = NRTYGeo::TGeoIndexFormatter::FormatFileName(context.DstDir.PathName(), "index");
            NRTYGeo::TGeoIndexFormatter::Save(outputPath, index);
            return true;
        }
    };
}
