#include <util/draft/datetime.h>
#include <util/stream/zlib.h>
#include <util/generic/buffer.h>
#include <util/stream/buffer.h>
#include <util/charset/wide.h>

#include <library/cpp/getopt/last_getopt.h>
#include <mapreduce/yt/interface/client.h>

#include <mapreduce/yt/common/config.h>
#include <wmconsole/version3/wmcutil/yt/yt_runner.h>
#include <wmconsole/version3/wmcutil/log.h>
#include <wmconsole/version3/wmcutil/args.h>


#include "compressor.h"
#include "util.h"
#include "run_config.h"

namespace NWebmaster {

static const size_t MAX_BUFFER_SIZE = 50 * 1024 * 1024;

void TPartitionAndOrderAwareReducer::Flush(ui32 tableId, const TString &partitionKey, const TString &orderKey,
    THolder<NUtils::TBufferedGZStream> &stream, TWriter *writer, bool isCombiner) {

    if (stream->Size() > 0) {
        stream->Finish();
        NYT::TNode resultRow;
        resultRow["TableId"] = tableId;
        resultRow["PartitionKey"] = partitionKey;
        resultRow["OrderKey"] = orderKey;
        resultRow["IsPacked"] = true;
        resultRow["data"] = TString(stream->Data(), stream->Size());
        writer->AddRow(resultRow, isCombiner ? 0 : tableId);
        stream.Reset(new NUtils::TBufferedGZStream());
    }
}

void TPartitionAndOrderAwareReducer::Do(TReader *reader, TWriter *writer) {
    const TRunConfig &config = TRunConfig::CInstance();
    bool isCombiner = config.Stage == "combine_v2";

    THolder<NUtils::TBufferedGZStream> gzStream(new NUtils::TBufferedGZStream());
    ui32 tableId = reader->GetRow()["TableId"].AsUint64();
    TString partitionKey = reader->GetRow()["PartitionKey"].AsString();
    TString orderKey;
    for (; reader->IsValid(); reader->Next()) {
        const NYT::TNode &row = reader->GetRow();
        if (orderKey.Empty()) {
            orderKey = row["OrderKey"].AsString();
        }
        if (row["IsPacked"].AsBool()) {
            writer->AddRow(row, isCombiner ? 0 : tableId);
            orderKey.clear();
        } else {
            const TString &data = row["data"].AsString();
            gzStream->Write(data.c_str(), data.size());
            if (gzStream->Size() >= MAX_BUFFER_SIZE) {
                Flush(tableId, partitionKey, orderKey, gzStream, writer, isCombiner);
                orderKey.clear();
            }
        }
    }
    Flush(tableId, partitionKey, orderKey, gzStream, writer, isCombiner);
}

} //namespace NWebmaster
