#pragma once

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

#include <util/folder/path.h>
#include <util/system/mutex.h>

namespace NRTYGeo {
    class TGeoIndexTree;

    using TPartKey = std::pair</*stream_id*/ui8, /*TKeyPrefix*/ui64>;

    class TGeoIndex {
    public:
        THashMap<TPartKey, THolder<TGeoIndexTree>> Trees;

    public:
        const TGeoIndexTree* GetTree(ui8 streamId, ui64 kps) const {
            auto i = Trees.find(TPartKey{streamId, kps});
            if (i != Trees.end())
                return i->second.Get();
            return nullptr;
        }

        size_t GetHitsCount() const {
            size_t result = 0;
            for (auto&& [_, pTree] : Trees) {
                result += pTree->GetData().size();
            }
            return result;
        }
    };

    class TGeoIndexBuilder {
    private:
        using TTrees = THashMap<TPartKey, THolder<TGeoIndexTreeBuilder>>;
        TTrees Trees;
    public:
        TGeoIndexTreeBuilder& GetChildBuilder(ui8 streamId, ui64 kps) {
            TPartKey key{streamId, kps};
            THolder<TGeoIndexTreeBuilder>& vl = Trees[key];
            if (!vl)
                vl = MakeHolder<TGeoIndexTreeBuilder>();
            return *vl;
        }

        void AddDoc(ui8 streamId, ui64 kps, ui32 docId, const NGeo::TGeoWindow& docRange) {
            if (Y_UNLIKELY(!docRange.IsValid()))
                return;
            TGeoIndexTreeBuilder& bld = GetChildBuilder(streamId, kps);
            if (Y_LIKELY(TGeoIndexTree::IsNormalized(docRange)))
                bld.AddHit(docId, docRange);
            else {
                TVector<NGeo::TGeoWindow> ranges = TGeoIndexTree::Normalize(docRange);
                for (const auto& range: ranges) {
                    bld.AddHit(docId, range);
                }
            }
        }

        void AddHit(ui8 streamId, ui64 kps, ui32 docId, const TGeoIndexTree::TKey& item) {
            GetChildBuilder(streamId, kps).AddHit(docId, item);
        }

        void Remap(const TVector<ui32>& remapTable) {
            for (auto&& [_, tree]: Trees) {
                tree->Remap(remapTable);
            }
        }

        void Finalize(TGeoIndex& dest) {
            for (std::pair<const TPartKey, THolder<TGeoIndexTreeBuilder>>& kv: Trees) {
                THolder<TGeoIndexTreeBuilder>& subTree = kv.second;
                if (!subTree->ReadOnly) {
                    subTree->Finalize();
                }
                THolder<TGeoIndexTree> subTreeFinal = MakeHolder<TGeoIndexTree>();
                subTree->AssignTo(*subTreeFinal);
                dest.Trees[kv.first] = std::move(subTreeFinal);
            }
            Trees.clear();
        }
    };

    class TGeoIndexFormatter {
    public:
        static constexpr TStringBuf Suffix = "geo.rty";

    public:
        static TFsPath FormatFileName(const TFsPath& dirName, const TString& indexPrefix) {
            return dirName / (indexPrefix + Suffix);  // "index_0000000000_0000002313/indexgeo.rty"
        }
        static TFsPath FormatFileName(const TPathName& path) {
            return FormatFileName(TFsPath(path.PathName()).Parent(), path.BaseName());
        }

        static void Save(const TFsPath& dest, const TGeoIndex& index);
        static void Load(const TFsPath& src, TGeoIndex& index);
        static bool Check(const TFsPath& src);
    };
}
