#pragma once

#include <kernel/doom/offroad_wad/offroad_wad_writer.h>
#include <saas/rtyserver/components/oxy/doom/offroad_wad_doc_buffer.h>

namespace NDoom {

/**
 * Optimized writer for merging offroad keyinv wad files.
 *
 * Takes pre-sorted hits, already groupped by DocId
 */

//TODO: (yrum) Remove copypaste - this should be moved to kernel/doom and be inherited from TOffroadWadWriter

template<EWadIndexType indexType, class Hit, class Vectorizer, class Subtractor, class PrefixVectorizer, class KeyPrefixGetter>
class TOffroadWadDocWriter {
private:
    using TKeyWriter = NOffroad::TFatKeyWriter<
        NOffroad::TTransformingKeyWriter<
            KeyPrefixGetter,
            NOffroad::TFatOffsetDataWriter<ui32, NOffroad::TNullSerializer>
        >,
        NOffroad::TKeyWriter<ui32, NOffroad::TUi32Vectorizer, NOffroad::TI1Subtractor>
    >;
    using THitWriter = NOffroad::TAdaptiveTupleWriter<Hit, Vectorizer, Subtractor, PrefixVectorizer>;
    using THitBuffer = TOffroadWadDocBuffer<Hit>;

public:
    using THit = Hit;
    using TKey = TString;
    using TKeyRef = TStringBuf;

    using THitTable = typename THitWriter::TTable;
    using TKeyTable = typename TKeyWriter::TTable;
    using THitModel = typename THitWriter::TModel;
    using TKeyModel = typename TKeyWriter::TModel;
    using TKeyBuffer = TOffroadKeyBuffer<TKey>;

    static_assert(THitWriter::Stages == 1, "Expected 1 stage for hit writer");
    static_assert(TKeyWriter::Stages == 1, "Expected 1 stage for key writer");

    enum {
        Stages = 1
    };

    TOffroadWadDocWriter() {}

    TOffroadWadDocWriter(const THitModel& hitModel, const TKeyModel& keyModel, const TString& path) {
        Reset(hitModel, keyModel, path);
    }

    TOffroadWadDocWriter(const THitModel& hitModel, const TKeyModel& keyModel, IOutputStream* output) {
        Reset(hitModel, keyModel, output);
    }

    void Reset(const THitModel& hitModel, const TKeyModel& keyModel, const TString& path) {
        Output_.Reset(new TOFStream(path));
        Reset(hitModel, keyModel, Output_.Get());
    }

    void Reset(const THitModel& hitModel, const TKeyModel& keyModel, IOutputStream* output) {
        HitTable_->Reset(hitModel);
        KeyTable_->Reset(keyModel);
        WadOutput_.Reset(output);

        /* Register wad types. */
        WadOutput_.RegisterDocLumpType(TWadLumpId(indexType, EWadLumpRole::Hits));
        WadOutput_.RegisterDocLumpType(TWadLumpId(indexType, EWadLumpRole::HitSub));

        /* Write out models right away. */
        IOutputStream* hitModelOutput = WadOutput_.StartGlobalLump(TWadLumpId(indexType, EWadLumpRole::HitsModel));
        hitModel.Save(hitModelOutput);

        IOutputStream* keyModelOutput = WadOutput_.StartGlobalLump(TWadLumpId(indexType, EWadLumpRole::KeysModel));
        keyModel.Save(keyModelOutput);

        /* Set up key i/o, key data will be written next. */
        FatOutput_.Reset();
        FatSubOutput_.Reset();
        IOutputStream* keyOutput = WadOutput_.StartGlobalLump(TWadLumpId(indexType, EWadLumpRole::Keys));
        KeyWriter_.Reset(&FatOutput_, &FatSubOutput_, KeyTable_.Get(), keyOutput);

        HitBuffer_.Reset();
        KeyBuffer_.Reset();
    }

    void WriteDoc(ui32 finalDocId) {
        HitBuffer_.Seek(finalDocId);
    }

    void WriteHit(const THit& hit) {
        HitBuffer_.Write(hit);
    }

    void WriteKey(const TKeyRef& key, ui32 hitsCount) {
        for(; hitsCount; --hitsCount)
            KeyBuffer_.AddHit(); //TODO: (yrum) AddHits(hitsCount)
        KeyBuffer_.AddKey(TString{key}); // This generates new sequential TermId
    }

    void Finish() {
        if (IsFinished())
            return;

        KeyBuffer_.PrepareKeys(true);
        HitBuffer_.RemapTermIds(KeyBuffer_.KeysMapping());
        TStringBuf lastKey;

        for (ui32 i = 0; i < KeyBuffer_.Keys().size(); ++i) {
            if (i > 0) {
                Y_ENSURE(lastKey < KeyBuffer_.Keys()[i]);
            }
            KeyWriter_.WriteKey(KeyBuffer_.Keys()[i], KeyBuffer_.SortedKeysMapping()[i]);
            lastKey = KeyBuffer_.Keys()[i];
        }

        KeyWriter_.Finish();

        IOutputStream* keyFatOutput = WadOutput_.StartGlobalLump(TWadLumpId(indexType, EWadLumpRole::KeyFat));
        FatOutput_.Flush(keyFatOutput);

        IOutputStream* keyIdxOutput = WadOutput_.StartGlobalLump(TWadLumpId(indexType, EWadLumpRole::KeyIdx));
        FatSubOutput_.Flush(keyIdxOutput);

        THitWriter hitWriter;
        for (size_t docId = 0; docId < HitBuffer_.Size(); docId++) {
            auto hits = HitBuffer_.Hits(docId);

            if (!hits.empty()) {
                hitWriter.Reset(
                    HitTable_.Get(),
                    WadOutput_.StartDocLump(docId, TWadLumpId(indexType, EWadLumpRole::Hits)),
                    WadOutput_.StartDocLump(docId, TWadLumpId(indexType, EWadLumpRole::HitSub))
                );
                for (const THit& hit : hits)
                    hitWriter.WriteHit(hit);
                hitWriter.Finish();
            }
        }
        WadOutput_.Finish();

        Output_.Reset();
    }

    bool IsFinished() const {
        return KeyWriter_.IsFinished();
    }

private:
    THolder<IOutputStream> Output_;
    THolder<THitTable> HitTable_ = MakeHolder<THitTable>();
    THolder<TKeyTable> KeyTable_ = MakeHolder<TKeyTable>();
    TMegaWadWriter WadOutput_;

    TAccumulatingOutput FatOutput_;
    TAccumulatingOutput FatSubOutput_;

    TKeyWriter KeyWriter_;
    THitBuffer HitBuffer_;
    TKeyBuffer KeyBuffer_;
};

} // namespace NDoom
