#include "memory_portions.h"

#include <kernel/walrus/advmerger.h>
#include <util/stream/file.h>

TMemoryPortionUsage::TMemoryPortionUsage(IYndexStorage::FORMAT format, int maxDocs, TString prefix, TString suffix, TString dir)
{
    Dir = dir;
    Format = format;
    Counter = 0;
    MaxDocs = maxDocs;
    PortionNum = 0;
    Prefix = prefix;
    Suffix = suffix;
    Portions.push_back(MakeHolder<NIndexerCore::TMemoryPortion>(Format));
}

bool TMemoryPortionUsage::IncDoc() {
    bool result = false;
    Counter++;
    if (Counter == MaxDocs) {
        StoreResult();
        result = true;
    }
    Portions.push_back(MakeHolder<NIndexerCore::TMemoryPortion>(Format));
    return result;
}

bool TMemoryPortionUsage::StoreResult() {
    if (Counter == 0)
        return false;
    const TString portionName(TString::Join(Prefix, ToString(PortionNum), 'p', Suffix));
    TAtomicSharedPtr<NIndexerCore::TMemoryPortion> mp = GetResultAndFlush();
    mp->Flush();
    TFixedBufferFileOutput foKeys(TString::Join(Dir, '/', portionName, 'k'));
    foKeys.Write(mp->GetKeyBuffer().Data(), mp->GetKeyBuffer().Size());
    TFixedBufferFileOutput foInvs(TString::Join(Dir, '/', portionName, 'i'));
    foInvs.Write(mp->GetInvBuffer().Data(), mp->GetInvBuffer().Size());

    PortionNum++;
    Counter = 0;
    return true;
}

TAtomicSharedPtr<NIndexerCore::TMemoryPortion> TMemoryPortionUsage::GetResultAndFlush() {
    TVector<TPortionBuffers> pbufs;

    for (size_t i = 0; i < Portions.size(); i++) {
        Portions[i]->Flush();
        pbufs.push_back(TPortionBuffers(
            Portions[i]->GetKeyBuffer().Data(),
            Portions[i]->GetKeyBuffer().Size(),
            Portions[i]->GetInvBuffer().Data(),
            Portions[i]->GetInvBuffer().Size()));
    }

    TAtomicSharedPtr<NIndexerCore::TMemoryPortion> result = new NIndexerCore::TMemoryPortion(Format);
    MergeMemoryPortions(&pbufs[0], Portions.size(), Format, nullptr, false, *result);
    Portions.clear();
    return result;
}
