#include "wad_merge.h"

#include "docid_remapper.h"

#include <robot/jupiter/library/rtdoc/file/docidmap_io.h>
#include <util/generic/set.h>
#include <util/stream/output.h>
#include <util/stream/file.h>

namespace NRtDoc {
    struct TWadMerger::TChunk {
        TChunk() {
        }

        TChunk(const TString& wadPath, THolder<TDocIdMap> docIdMap, IWadPatcher::TPtr wadPatcher, const TString& filename) {
            WadPath = wadPath;
            Wad = IWad::Open(wadPath, false);

            Y_ENSURE(docIdMap);

            DocIdMap.Swap(docIdMap);
            WadPatcher = wadPatcher;

            DocLumps = Wad->DocLumps();
            GlobalLumps = Wad->GlobalLumps();
            Name = filename;
            Mapping.resize(DocLumps.size());
            Wad->MapDocLumps(DocLumps, Mapping);
        }

        THolder<TDocIdMap> DocIdMap;
        IWadPatcher::TPtr WadPatcher;

        TString Name;
        TString WadPath;
        THolder<IWad> Wad = nullptr;
        TVector<TWadLumpId> DocLumps;
        TVector<TWadLumpId> GlobalLumps;
        TVector<size_t> Mapping;
    };

    struct TWadMerger::TCheckContext {
        TMap<TWadLumpId, TBlob> GlobalLumps;
        TSet<TWadLumpId> DocLumpTypes;
    };

    TWadMerger::TWadMerger() {
    }

    TWadMerger::TWadMerger(const TString& fileName, const TString& mappingFileName)
        : MappingFileName(mappingFileName) //TODO(yrum): MappingFileName is not needed
    {
        Init(fileName);
    }

    void TWadMerger::Init(const TString& fileName) {
        FileName = fileName;
    }

    TWadMerger::~TWadMerger() {
    }

    void TWadMerger::Add(const TString& wadPath, THolder<TDocIdMap> docIdMap, IWadPatcher::TPtr wadPatcher, const TString& name, bool /*isDelta*/) {
        Y_ENSURE(!CheckCtx_);
        Chunks_.emplace_back(wadPath, std::move(docIdMap), wadPatcher, name);
    }

    inline std::pair<bool, bool> CmpBlob(const TBlob& r1, const TBlob& r2) {
        std::pair<bool, bool> res(false, false);
        res.first = (r1.Size() == r2.Size());
        if (res.first)
            res.second = (memcmp(r1.Data(), r2.Data(), r1.Size()) == 0);
        return res;
    }

    static TString DumpWadIds(const TVector<TWadLumpId>& ids) {
        TStringStream s;
        for (auto id : ids) {
            if (!s.Empty())
                s << " ";
            s << id;
        }
        return s.Str();
    }

    template <typename TData>
    static TString DumpWadIds(const TMap<TWadLumpId, TData>& idsKeys) {
        TVector<TWadLumpId> ids;
        for (auto kv : idsKeys) {
            ids.push_back(kv.first);
        }
        return DumpWadIds(ids);
    }

    template <typename T1, typename T2>
    static TString DiffWadIds(const char kind[], const TString& aName, const T1& a, const TString& bName, const T2& b) {
        TStringStream s;
        s << kind << " differ between " << aName << " and " << bName << ": [" << DumpWadIds(a) << "] != [" << DumpWadIds(b) << "]";
        return s.Str();
    }

    void TWadMerger::CheckMergeable() {
        CheckCtx_.Destroy();

        auto ctx = MakeHolder<TCheckContext>();

        Y_ENSURE(!Chunks_.empty(), "No input files");
        if (Chunks_.size() == 1)
            return;

        //TODO(yrum): ensure that headChunk is built with the last builder version
        const auto* headChunk = &Chunks_.back(); // the chunk with the most actual metadata (one that is built last)
        for (TWadLumpId type : headChunk->GlobalLumps) {
            TBlob blob = headChunk->Wad->LoadGlobalLump(type);
            auto rc = ctx->GlobalLumps.insert(std::make_pair(type, blob));
            if (!rc.second) {
                blob.Drop();
            }
            Y_ENSURE(rc.second);
        }
        for (TWadLumpId type : headChunk->DocLumps) {
            Y_ENSURE(!ctx->DocLumpTypes.contains(type));
            ctx->DocLumpTypes.insert(type);
        }

        auto chunk = headChunk;
        for (++chunk; chunk != Chunks_.end(); ++chunk) {
            Y_ENSURE(chunk->GlobalLumps.size() == ctx->GlobalLumps.size(), DiffWadIds("GlobalLumps", headChunk->Name, headChunk->GlobalLumps, chunk->Name, chunk->GlobalLumps));
            for (TWadLumpId type : headChunk->GlobalLumps) {
                TBlob blob = chunk->Wad->LoadGlobalLump(type);
                const bool exists = ctx->GlobalLumps.contains(type);
                Y_ENSURE(exists, "GlobalLumps" << "differ between " << headChunk->Name << " and " << chunk->Name << ": missing " << type);
                const bool equal = exists && CmpBlob(blob, ctx->GlobalLumps.at(type)).second;
                Y_ENSURE(equal || type.Role == EWadLumpRole::StructSize, "GlobalLumps" << "differ between " << headChunk->Name << " and " << chunk->Name << ": " << type);
            };

            Y_ENSURE(chunk->DocLumps.size() == ctx->DocLumpTypes.size(), DiffWadIds("DocLumps", headChunk->Name, headChunk->GlobalLumps, chunk->Name, chunk->GlobalLumps));
            for (TWadLumpId type : headChunk->DocLumps) {
                Y_ENSURE(ctx->DocLumpTypes.contains(type), "DocLumps" << "differ between " << headChunk->Name << " and " << chunk->Name << ": missing" << type);
            };
        }

        CheckCtx_.Swap(ctx);
    }

    void TWadMerger::Finish() {
        Y_ENSURE(FileName);
        Y_ENSURE(!Chunks_.empty(), "No input files");

        if (!CheckCtx_) {
            CheckMergeable(); //throws an exception if the Wads are not compatible
        }

        /* Calc doc count. */
        TDocIdRemapper docidRemapper(MappingFileName, Chunks_.size());
        for (const TChunk& chunk : Chunks_) {
            docidRemapper.AddSegmentSize(chunk.Wad->Size());
        }

        TMegaWadWriter writer(FileName);

        // note: CheckMergable() should be called to avoid incorrect Lumps formation
        const TChunk& headChunk = Chunks_.back();

        /* Note that we're registering these in the same order as they appeared in source wads. */
        THashSet<TWadLumpId> knownWads;
        for (TWadLumpId type : headChunk.Wad->DocLumps()) {
            writer.RegisterDocLumpType(type);
            knownWads.insert(type);
        }

        /* Save global lumps. */
        for (TWadLumpId type : headChunk.GlobalLumps) {
            TBlob blob = headChunk.Wad->LoadGlobalLump(type);
            writer.WriteGlobalLump(type, TArrayRef<const char>(blob.AsCharPtr(), blob.Size()));
        }

        /* Save doc lumps. */
        docidRemapper.InitFinalDocIds();

        TMap<ui32, TVector<std::pair<const NDoom::TWadLumpId, TArrayRef<const char>>>> sortedLumps;

        for (const TChunk& chunk : Chunks_) {
            if (chunk.WadPatcher) {
                chunk.WadPatcher->Reset(chunk.Wad.Get()); // read offroad model and struct size
            }
            docidRemapper.ProcessSegment(*chunk.DocIdMap, [&](const ui32 docId, const ui32 outputDocId) {
                TVector<TArrayRef<const char>> regs;
                regs.resize(chunk.Mapping.size());
                TBlob temp = chunk.Wad->LoadDocLumps(docId, chunk.Mapping, regs);
                if (IsSortNeeded() && !regs.empty()) {
                    sortedLumps[outputDocId].clear();
                }

                for (size_t i = 0; i < regs.size(); i++) {
                    const auto lumpType = chunk.DocLumps[i];
                    if (!knownWads.contains(lumpType)) {
                        continue;
                    }
                    const TArrayRef<const char>& srcLump = regs[i];
                    TArrayRef<const char> resultLump = chunk.WadPatcher ? chunk.WadPatcher->ProcessDoc(outputDocId, lumpType, srcLump) : srcLump;
                    if (!IsSortNeeded()) {
                        writer.WriteDocLump(outputDocId, lumpType, resultLump);
                    } else {
                        // need to be sorted first
                        sortedLumps[outputDocId].push_back({lumpType, resultLump});
                    }
                }
            });
        }

        if (IsSortNeeded()) {
            for (const auto& [id, lumps] : sortedLumps) {
                for (const auto& [lumpType, lumpData] : lumps) {
                    writer.WriteDocLump(id, lumpType, lumpData);
                }
            }
        }
        writer.Finish();

        docidRemapper.Finish();
    }
}
