#include "yt_file_fetcher.h"

#include <saas/library/yt/stream/yt_chunked_input.h>

#include <mapreduce/yt/common/config.h>
#include <mapreduce/yt/interface/client.h>
#include <mapreduce/yt/interface/io.h>

#include <util/string/cast.h>

TYTFileFetcher::TYTFileFetcher(const TString& ytProxy, const TString& ytToken, const TString& ytHosts)
    : ClientPtr(NYT::CreateClient(ytProxy, NYT::TCreateClientOptions().Token(ytToken)))
{
    Y_ENSURE(!!ClientPtr);
    if (!ytHosts.empty()) {
        NYT::TConfig::Get()->Hosts = ytHosts;
    }
}

TYTFileFetcher::~TYTFileFetcher() = default;

void TYTFileFetcher::Fetch(NYT::TRichYPath ytPath, TFsPath outputRoot, ui32 writeBytesPerSec) {
    EYTBlobFormat format = GetBlobFormatFromAttrs(ClientPtr, ytPath);
    NYT::TTableReaderPtr<NYT::Message> reader;
    if (format == EYTBlobFormat::Old) {
        reader = NYT::CreateGenericProtobufReader(ClientPtr->CreateTableReader<NSaas::TYTBlobBase>(ytPath));
    } else {
        reader = NYT::CreateGenericProtobufReader(ClientPtr->CreateTableReader<NSaas::TYTNewBlobBase>(ytPath));
    }
    if (writeBytesPerSec > 0) {
        TYTChunkedInputToThrottledFiles in(reader.Get(), writeBytesPerSec, outputRoot, format);
        in.ProcessAll();
    } else {
        TYTChunkedInputToFiles in(reader.Get(), outputRoot, format);
        in.ProcessAll();
    }
}
