#include "writer.h"

#include <library/cpp/logger/global/global.h>

namespace NRTYAnn {

    TStreamDataWriter::TStreamDataWriter(const TArchiveOwner::TPtr& storage,
                                         TAtomicSharedPtr<IDataBlockAccessor> accessor)
    {
        DocumentsStorage = storage;
        DataAccessor = accessor;
        DocId = Max<ui32>();
    }

    void TStreamDataWriter::Add(ui32 breakId, ui32 regionId, ui32 streamId, TArrayRef<const char> data) {
        TFixedSizeKey<3> key;
        key.Set(0, breakId).Set(1, regionId).Set(2, 0);
        ui32 offset = StreamData->Size();
        if (!IndexBuilder->Find(key.data(), key.size(), &offset)) {
            StreamData->Advance(DataAccessor->GetSize());
            memset(StreamData->Data() + offset, 0, DataAccessor->GetSize());  // potentially invalid zero fill, which affects docs relevance: SAASSUP-4278
            IndexBuilder->Add(key.data(), key.size(), offset);
        }
        DataAccessor->SetBlockBegining(StreamData->Data() + offset);
        DataAccessor->Set(streamId, data);
    };

    TBlob TStreamDataWriter::GetFinalBlob() const {
        TBuffer buf;
        buf.Resize(sizeof(TStreamDataHeader));

        TBufferOutput out(buf);
        TStreamDataHeader header;
        header.Size = IndexBuilder->Save(out);
        *((TStreamDataHeader*)buf.Data()) = header;
        buf.Append(StreamData->Data(), StreamData->Size());
        return TBlob::FromBuffer(buf);
    }

    void TStreamDataWriter::StartDoc(ui32 docId) {
        DocId = docId;
        StreamData.Reset(new TBuffer);
        IndexBuilder.Reset(new TCompactTrieBuilder<ui32, ui32>);
    }

    void TStreamDataWriter::FinishDoc() {
        TBlob data = GetFinalBlob();
        if (data.Size() == sizeof(TStreamDataHeader))
            return;
        DocumentsStorage->PutDocument(data, DocId);
    }
}
