#pragma once

#include <robot/jupiter/library/rtdoc/merger/wad_merger.h>
#include <kernel/doom/wad/mega_wad.h>
#include <kernel/doom/offroad_reg_herf_wad/reg_herf_io.h>
#include <kernel/doom/standard_models_storage/standard_models_storage.h>
#include <util/generic/ptr.h>

namespace NRtDoc {
    using namespace NDoom;

    //
    // Base IWadMerger implementation for TWadDecodedMerger
    //
    class TWadDecodedMergerBase : public IWadMerger {
    protected:
        struct TChunk {
        public:
            TChunk(const TString& wadPath, THolder<TDocIdMap> docIdMap, const TString& filename) {
                Y_ENSURE(wadPath);
                WadPath = wadPath;
                Wad = IWad::Open(wadPath, false);

                Y_ENSURE(docIdMap);
                Y_ENSURE(docIdMap->GetData()->size() >= Wad->Size());

                DocIdMap.Swap(docIdMap);
                Name = filename;
            }

            THolder<TDocIdMap> DocIdMap;
            TString Name;
            TString WadPath;
            THolder<IWad> Wad = nullptr;
        };

    public:
        TWadDecodedMergerBase();

        ~TWadDecodedMergerBase();

        virtual void Init(const TString& fileName) override;

        virtual void Add(const TString& wadPath, THolder<TDocIdMap> docIdMap, IWadPatcher::TPtr wadPatcher, const TString& name, bool /*isDelta*/) override;

        virtual void CheckMergeable() override {
            // no checks
        };


    public:
        virtual bool Empty() const override {
            return Chunks_.empty();
        }

    protected:
        TVector<TChunk> Chunks_;
        TString FileName;
    };

    //
    // Templates for particular Wad IOs
    //
    template <typename TIo, typename TKeyAccessor>
    class TWadDecodedDataStorage {
    public:
        using TKey = typename TIo::TKey;
        using TData = typename TIo::TData;
        struct TRecord {
            TKey Key;
            TBlob Value;
        };

        using TVec = TVector<TRecord>;
        using TIterator = typename TVec::const_iterator;

    private:
        THashMap<ui32, TVec> Lookup;
        ui32 LastDocId = Max<ui32>();
    public:
        void Add(const TKey& key, const TData& data) {
            const ui32 docId = TKeyAccessor::DocId(key);
            if (LastDocId != docId) {
                if (Lookup.contains(docId))
                    Lookup[docId].clear(); // to replace data for same docId
                LastDocId = docId;
            }

            Lookup[docId].emplace_back(TRecord{key, TBlob::Copy(&data, sizeof(data))});
        }

        std::pair<TIterator, TIterator> GetRange(ui32 docId) const {
            auto lookup = Lookup.find(docId);
            if (Y_UNLIKELY(lookup == Lookup.end())) {
                return std::make_pair(TIterator(), TIterator());
            }
            return std::make_pair(lookup->second.cbegin(), lookup->second.cend());
        }
    };


    template <typename TIo, typename TKeyAccessor>
    class TWadDecodedMerger : public TWadDecodedMergerBase {
    private:
        TWadDecodedDataStorage<TIo, TKeyAccessor> Data_;

    public:
        //! Process all IWads and write merged WAD.
        virtual void Finish() override {
            using TKey = typename TIo::TKey;
            using TData = typename TIo::TData;
            using TReader = typename TIo::TReader;
            using TWriter = typename TIo::TWriter;
            using TModel = typename TIo::TWriter::TModel;

            ui32 maxId = Max<ui32>();

            for (const auto& s : Chunks_) {
                TReader rd(s.Wad.Get());
                TKey key;
                const TData* data;
                while (rd.Next(&key, &data)) {
                    if (!data) {
                        continue;
                    }

                    ui32 oldId = TKeyAccessor::DocId(key);
                    ui32 newId = s.DocIdMap->Map(oldId);
                    if (newId == TDocIdMap::DeletedDocument()) {
                        continue;
                    }

                    TKey newKey = TKeyAccessor::SetDocId(key, newId);
                    Data_.Add(newKey, *data);

                    if (maxId < newId || maxId == Max<ui32>()) {
                        maxId = newId;
                    }
                }
            }

            {
                TModel model = TStandardIoModelsStorage::Model<TModel>(TIo::DefaultModel);
                TWriter wr(model, FileName);
                if (maxId != Max<ui32>()) {
                    for (ui32 id = 0; id <= maxId; ++id) {
                        auto range = Data_.GetRange(id);
                        for (auto i = range.first; i != range.second; ++i) {
                            wr.Write(i->Key, static_cast<const TData*>(i->Value.Data()));
                        }
                    }
                }
                wr.Finish();
            }
        }
    };
}
