#include "writer.h"

#include <saas/rtyserver/common/should_stop.h>

#include <ysite/directtext/sentences_lens/sentenceslensproc.h>
#include <saas/rtyserver/merger/interface/common.h>

TRTYSentenceLengthWriter::TRTYSentenceLengthWriter(const TFsPath& resultPrefix, const TFsPath& tempPrefix)
    : TempFile(tempPrefix.GetPath() + "sent")
    , FinalFile(resultPrefix.GetPath() + "sent")
{
    Writer.Reset(new TSentenceLengthsWriter(TempFile));
}

void TRTYSentenceLengthWriter::AddPacked(ui32 docId, const TSentenceLengths& lengthsPacked) {
    TGuard<TMutex> guard(Mutex);
    Y_ASSERT(Writer);
    Writer->AddPacked(docId, lengthsPacked);
}

void TRTYSentenceLengthWriter::Add(ui32 docId, const ui8* begin, size_t size) {
    TGuard<TMutex> guard(Mutex);
    Y_ASSERT(Writer);
    Writer->Add(docId, begin, size);
}

void TRTYSentenceLengthWriter::LoadLength(ui32 docId, const TSentenceLensProcessor& proc) {
    TGuard<TMutex> guard(Mutex);
    TBufferOutput sentenceLengths;
    proc.WriteWithZeroSentence(&sentenceLengths);
    Writer->Add(docId, (const ui8*)sentenceLengths.Buffer().Data(), sentenceLengths.Buffer().Size());
}

void TRTYSentenceLengthWriter::Close(const TVector<ui32>* remapTable, const std::atomic<bool> *rigidStopSignal) {
    Writer.Reset(nullptr);

    if (!remapTable) {
        TempFile.RenameTo(FinalFile);
        return;
    }

    const size_t remapSize = remapTable->size();
    {
        TSentenceLengthsReader reader(TempFile);
        TSentenceLengthsWriter writer(FinalFile);
        TSentenceLengths lengths;
        for (ui32 docId = 0; docId < reader.GetSize() && !ShouldStop(rigidStopSignal); ++docId) {
            reader.Get(docId, &lengths);
            const ui32 finalDocId = docId < remapSize ? remapTable->at(docId) : docId;
            if (finalDocId == REMAP_NOWHERE)
                continue;

            writer.Add(finalDocId, lengths);
        }
    }
    TempFile.DeleteIfExists();
}
