#include <util/generic/string.h>
#include <util/string/builder.h>
#include <util/string/join.h>
#include <util/string/printf.h>

#include "yt_runner.h"

namespace NWebmaster {
namespace NYTRunner {

void (*TOpRunner::LogInfo)(const TString &msg) = &TOpRunner::DefaultLogInfo;

TOpRunner::TOpRunner(NYT::IClientBasePtr client)
    : Client(client)
{
}

TOpRunner& TOpRunner::Archive(const TString& table, int ctx, size_t desiredChunkCount, const TString& eCodec, const TString& cCodec) {
    PrepareParams();
    return RunTask(new TArchiveDefragTask(Client, table, desiredChunkCount, eCodec, cCodec), ctx);
}

TOpRunner& TOpRunner::Defrag(const TString& table, int ctx, size_t desiredChunkCount) {
    PrepareParams();
    return RunTask(new TArchiveDefragTask(Client, table, desiredChunkCount), ctx);
}

TOpRunner& TOpRunner::Comment(const TString &comment) {
    CommentStr = comment;
    return *this;
}

void TOpRunner::CommentLogInfo(TString msg) const {
    LogInfo(GetCommentedMsg(msg));
}

TOpRunner& TOpRunner::Copy(const TString &src, const TString &dst) {
    TString opComment = GetCommentedMsg(TStringBuilder() << "copy operation " << src << " -> " << dst);
    TOpRunner::LogInfo(opComment);
    NYT::TCopyOptions opts;
    opts.Recursive(true);
    opts.Force(true);
    Client->Copy(src, dst, opts);
    TOpRunner::LogInfo(opComment + " - done");
    return *this;
}

void TOpRunner::DefaultLogInfo(const TString &msg) {
    Cout << msg << Endl;
}

TOpRunner& TOpRunner::Drop(const TString &table) {
    NYT::TRemoveOptions opts;
    opts.Recursive(true)
        .Force(true);
    Client->Remove(table, opts);
    CommentLogInfo("dropped table " + table);
    return *this;
}

TOpRunner& TOpRunner::DropByPrefix(const TString &prefix) {
    TDeque<NYTUtils::TTableInfo> tables;
    NYTUtils::GetTableList(Client, prefix, tables, Max<size_t>());
    for (const NYTUtils::TTableInfo &table : tables) {
        Drop(table.Name);
    }
    return *this;
}

TOpRunner& TOpRunner::EnableOrderedMap() {
    PrepareParams();
    Params->OrderedMap = true;
    return *this;
}

TOpRunner& TOpRunner::EnableLivePreview() {
    PrepareParams();
    Params->EnableLivePreviewFlag = true;
    return *this;
}

TString TOpRunner::GetCommentedMsg(const TString &msg) const {
    if (!CommentStr.empty()) {
        return CommentStr + " " + msg;
    }

    return msg;
}

TString TOpRunner::GetParametersComment(TParams &params) {
    TVector<TString> input, output;

    for (const auto &table : params.SrcTables) {
        input.push_back(table.RichPath.Path_);
    }

    for (const auto &table : params.DstTables) {
        output.push_back(table.RichPath.Path_);
    }

    if (input.size() > 5) {
        size_t total = input.size();
        input.resize(5);
        input.push_back(Sprintf("...%lu tables", total));
    }

    if (output.size() > 5) {
        size_t total = output.size();
        output.resize(5);
        output.push_back(Sprintf("...%lu tables", total));
    }

    return "[input: " + JoinSeq(", ", input) + "; output: " + JoinSeq(", ", output) + "]";
}

TOpRunner& TOpRunner::JoinBy(const NYT::TSortColumns &keys) {
    PrepareParams();
    Params->JoinBy.Reset(new NYT::TSortColumns(keys));
    return *this;
}

TOpRunner& TOpRunner::InputImpl(const NYT::TRichYPath &input, TOpRunner::TTable::EType type) {
    PrepareParams();

    CommentLogInfo("input " + input.Path_);
    Params->AddInputTable(input,  type);

    return *this;
}

TOpRunner& TOpRunner::InputByPrefixImpl(const NYT::TRichYPath &prefix, TOpRunner::TTable::EType type) {
    PrepareParams();

    NYT::TRichYPath pathTpl = prefix;
    TDeque<NYTUtils::TTableInfo> tables;
    NYTUtils::GetTableList(Client, pathTpl.Path_, tables, Max<size_t>());
    for (const NYTUtils::TTableInfo &table : tables) {
        pathTpl.Path_ = table.Name;
        CommentLogInfo("input " + table.Name);
        Params->AddInputTable(pathTpl,  type);
    }

    return *this;
}

TOpRunner& TOpRunner::Input(const NYT::TRichYPath &input) {
    return InputImpl(input, TTable::UNDEFINED);
}

TOpRunner& TOpRunner::InputNode(const NYT::TRichYPath &input) {
    return InputImpl(input, TTable::NODE);
}

TOpRunner& TOpRunner::InputProto(const NYT::TRichYPath &input) {
    return InputImpl(input, TTable::PROTO);
}

TOpRunner& TOpRunner::InputYaMR(const NYT::TRichYPath &input) {
    return InputImpl(input, TTable::YAMR);
}

TOpRunner& TOpRunner::InputByPrefix(const NYT::TRichYPath &prefix) {
    return InputByPrefixImpl(prefix, TTable::UNDEFINED);
}

TOpRunner& TOpRunner::InputNodeByPrefix(const NYT::TRichYPath &prefix) {
    return InputByPrefixImpl(prefix, TTable::NODE);
}

TOpRunner& TOpRunner::InputProtoByPrefix(const NYT::TRichYPath &prefix) {
    return InputByPrefixImpl(prefix, TTable::PROTO);
}

TOpRunner& TOpRunner::InputYaMRByPrefix(const NYT::TRichYPath &prefix) {
    return InputByPrefixImpl(prefix, TTable::YAMR);
}

TOpRunner& TOpRunner::File(const TString &file) {
    PrepareParams();
    Params->MapperUserJobSpec.AddFile(file);
    Params->CombinerUserJobSpec.AddFile(file);
    Params->ReducerUserJobSpec.AddFile(file);
    CommentLogInfo("file " + file);
    return *this;
}

TOpRunner& TOpRunner::LocalFile(const TString &file) {
    PrepareParams();
    Params->MapperUserJobSpec.AddLocalFile(file);
    Params->CombinerUserJobSpec.AddLocalFile(file);
    Params->ReducerUserJobSpec.AddLocalFile(file);
    CommentLogInfo("local file " + file);
    return *this;
}

TOpRunner& TOpRunner::MaxRowWeight(size_t bytes) {
    PrepareParams();
    NYT::TNode config;
    config["table_writer"]["max_row_weight"] = bytes;
    Spec("job_io", config);
    Spec("map_job_io", config);
    Spec("reduce_job_io", config);
    Spec("sort_job_io", config);
    Spec("partition_job_io", config);
    Spec("merge_job_io", config);
    return *this;
}

TOpRunner& TOpRunner::MemoryLimit(size_t limit) {
    PrepareParams();
    Params->MapperUserJobSpec.MemoryLimit(limit);
    Params->CombinerUserJobSpec.MemoryLimit(limit);
    Params->ReducerUserJobSpec.MemoryLimit(limit);
    return *this;
}

TOpRunner& TOpRunner::MapperMemoryLimit(size_t limit) {
    PrepareParams();
    Params->MapperUserJobSpec.MemoryLimit(limit);
    return *this;
}

TOpRunner& TOpRunner::CombinerMemoryLimit(size_t limit) {
    PrepareParams();
    Params->CombinerUserJobSpec.MemoryLimit(limit);
    return *this;
}

TOpRunner& TOpRunner::ReducerMemoryLimit(size_t limit) {
    PrepareParams();
    Params->ReducerUserJobSpec.MemoryLimit(limit);
    return *this;
}

TOpRunner& TOpRunner::DataSizePerJob(size_t dspj) {
    PrepareParams();
    Params->DataSizePerJob = dspj;
    return *this;
}

TOpRunner& TOpRunner::JobCount(size_t jc) {
    PrepareParams();
    Params->JobCount = jc;
    return *this;
}

TOpRunner& TOpRunner::PartitionCount(size_t pc) {
    PrepareParams();
    Params->PartitionCount = pc;
    return *this;
}

TOpRunner& TOpRunner::PartitionDataSize(size_t pds) {
    PrepareParams();
    Params->PartitionDataSize = pds;
    return *this;
}

TOpRunner& TOpRunner::OutputImpl(const NYT::TRichYPath &output, TOpRunner::TTable::EType type) {
    PrepareParams();
    CommentLogInfo("output " + output.Path_);
    Params->AddOutputTable(output, type);
    return *this;
}

TOpRunner& TOpRunner::Output(const NYT::TRichYPath &output) {
    return OutputImpl(output, TTable::UNDEFINED);
}

TOpRunner& TOpRunner::OutputNode(const NYT::TRichYPath &output) {
    return OutputImpl(output, TTable::NODE);
}

TOpRunner& TOpRunner::OutputProto(const NYT::TRichYPath &output) {
    return OutputImpl(output, TTable::PROTO);
}

TOpRunner& TOpRunner::OutputYaMR(const NYT::TRichYPath &output) {
    return OutputImpl(output, TTable::YAMR);
}

TOpRunner& TOpRunner::Merge(int ctx) {
    PrepareParams();
    TString opComment = GetCommentedMsg(TString("merge operation ") + GetParametersComment(*Params));
    return RunTask(new TMergeTask(Client, Params, opComment), ctx);
}

TOpRunner& TOpRunner::MergeBy(const NYT::TSortColumns &keys) {
    PrepareParams();
    Params->MergeBy.Reset(new NYT::TSortColumns(keys));
    return *this;
}

TOpRunner& TOpRunner::Move(const TString &src, const TString &dst) {
    TString opComment = GetCommentedMsg(TStringBuilder() << "move operation " << src << " -> " << dst);
    TOpRunner::LogInfo(opComment);
    NYT::TMoveOptions opts;
    opts.Recursive(true);
    opts.Force(true);
    Client->Move(src, dst, opts);
    TOpRunner::LogInfo(opComment + " - done");
    return *this;
}

TOpRunner& TOpRunner::ReduceBy(const NYT::TSortColumns &keys) {
    PrepareParams();
    Params->ReduceBy.Reset(new NYT::TSortColumns(keys));
    return *this;
}

TOpRunner& TOpRunner::Sort(const TString &src, const TString &dst, int ctx) {
    PrepareParams();

    if (!Params->SortBy) {
        ythrow yexception() << "sort: SortBy columns are not set";
    }

    return RunTask(new TSortTask(Client, Params, src, dst), ctx);
}

TOpRunner& TOpRunner::Sort(const TString &src, int ctx) {
    return Sort(src, src, ctx);
}

TOpRunner& TOpRunner::SortBy(const NYT::TSortColumns &keys) {
    PrepareParams();
    Params->SortBy.Reset(new NYT::TSortColumns(keys));
    return *this;
}

TOpRunner& TOpRunner::UseTmpfsInMapper() {
    NYT::TNode tmpfsSpec = NYT::TNode::CreateMap();
    tmpfsSpec["tmpfs_path"] = ".";
    tmpfsSpec["copy_files"] = true;
    Spec("mapper", tmpfsSpec);
    return *this;
}

TOpRunner& TOpRunner::UseTmpfsInReducer() {
    NYT::TNode tmpfsSpec = NYT::TNode::CreateMap();
    tmpfsSpec["tmpfs_path"] = ".";
    tmpfsSpec["copy_files"] = true;
    Spec("reducer", tmpfsSpec);
    return *this;
}

TOpRunner& TOpRunner::UseTmpfsInReduceCombiner() {
    NYT::TNode tmpfsSpec = NYT::TNode::CreateMap();
    tmpfsSpec["tmpfs_path"] = ".";
    tmpfsSpec["copy_files"] = true;
    Spec("reduce_combiner", tmpfsSpec);
    return *this;
}

TOpRunner& TOpRunner::UseTmpfs() {
    UseTmpfsInMapper();
    UseTmpfsInReduceCombiner();
    UseTmpfsInReducer();
    return *this;
}

TOpRunner& TOpRunner::Wait(size_t ctx) {
    CtxQueue[ctx].clear();
    return *this;
}

} //namespace NYTRunner
} //namespace NWebmaster
