#pragma once

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

#include <saas/rtyserver/components/l2/l2_manager.h>
#include <saas/rtyserver/model/index.h>
#include <library/cpp/logger/global/global.h>

namespace NRTYServer {
    class TMinGeoIndexData: public TNonCopyable {
    public:
        const TFsPath IndexDir;
        THolder<NRTYGeo::TGeoIndex> Inv;
        const IL2RawAccessor& Archive;

    public:
        TMinGeoIndexData(const TFsPath& indexDir, const IL2RawAccessor& archive)
            : IndexDir(indexDir)
            , Archive(archive)
        {
            Y_ENSURE(IndexDir.IsDirectory());
        }

        void Open() {
            const TFsPath filePath = NRTYGeo::TGeoIndexFormatter::FormatFileName(IndexDir, "index");
            if (filePath.Exists()) {
                THolder<NRTYGeo::TGeoIndex> inv = MakeHolder<NRTYGeo::TGeoIndex>();
                NRTYGeo::TGeoIndexFormatter::Load(filePath, *inv);
                Inv.Swap(inv);
            } else {
                WARNING_LOG << "File " << filePath.Basename() << " is not found, Geo functionality will be limited" << Endl;
            }
        }
    };

    class TMinGeoSearchContext: public TComponentSearcher {
    public:
        enum EPruningMode {
            pmNone,
            pmExpand,
            pmAuto
        };
        struct TPruningOpts {
            EPruningMode Mode = pmNone;
            ui32 DocCount = Max<ui32>();
            TMaybe<NGeo::TSize> SpanLimit;
        };
    private:
        const TMinGeoIndexData& IndexData;
        NGeo::TGeoWindow UserRequest;
        TPruningOpts PruningOpts;
        TVector<ui64> KeyPrefixes;

    public:
        TMinGeoSearchContext(const TMinGeoIndexData& idx)
            : IndexData(idx)
        {
        }

        bool IsUsed(const TRTYSearchRequestContext& context) const override {
            return nullptr != context.GetParameters(CgiCompName) && IndexData.Inv;
        }

        bool CanLookup() const override {
            return UserRequest.IsValid();
        }

        TAtomicSharedPtr<IDocProcessor> CreateFilter() override {
            // No extra filtration as for now (maybe added later)
            return nullptr;
        }

        // Parse CGI
        bool ParseContext(const TRTYSearchRequestContext& context, TMessagesCollector& errors) override;

        // Do Search
        TVector<ui32> Lookup() override {
            TVector<ui32> result;
            Y_VERIFY(!KeyPrefixes.empty(), "ParseContext was not called");
 
            if (Y_LIKELY(KeyPrefixes.size() == 1)) {
                // Single KPS case
                result = LookupImpl(KeyPrefixes.front());
            } else {
                // Multiple KPS case.
                // we support this, but do not recommend using it. It costs extra CPU and breaks some features.
                for (ui64 kps: KeyPrefixes) {
                    TVector<ui32> tmp = LookupImpl(kps);
                    result.insert(result.end(), tmp.begin(), tmp.end());
                }
                Sort(result);
                auto pEnd = std::unique(result.begin(), result.end());
                result.resize(std::distance(result.begin(), pEnd));
            }
            if (Y_UNLIKELY(result.empty())) {
                result = {Max<ui32>()}; // DocidFilter (TDocIdIterator) uses this convention to enforce zero result (see: AutoAcceptedIterator)
            }
            return result;
        }
    private:
        TVector<ui32> LookupImpl(ui64 keyPrefix) {
            TVector<ui32> result;
            constexpr ui8 streamId = 0; //TODO(SAAS-5949): allow user to specify streamId in CGI parameters
            const NRTYGeo::TGeoIndexTree* kdTree = IndexData.Inv->GetTree(streamId, keyPrefix);
            if (Y_LIKELY(kdTree)) {
                if (PruningOpts.Mode == pmNone || PruningOpts.DocCount == Max<ui32>()) {
                    result = kdTree->FindIntersections(UserRequest, Max<ui32>());
                } else {
                    Y_ASSERT(PruningOpts.Mode == pmExpand || PruningOpts.Mode == pmAuto);
                    result = kdTree->FindIntersectionsWithDynamicPruning(
                            UserRequest,
                            PruningOpts.DocCount,
                            /*mayCollapse=*/PruningOpts.Mode == pmAuto,
                            /*mayExpand=*/true,
                            /*maxSize=*/PruningOpts.SpanLimit.Get());
                }
            }
            return result;
        }


    public:
        static const TString CgiCompName;
    };
}
