#include <maps/libs/log8/include/log8.h>

#include <maps/wikimap/mapspro/services/autocart/pipeline/libs/yt_utils/include/op_wrapper.h>

#include <mapreduce/yt/interface/client.h>

namespace maps::wiki::autocart::pipeline {

namespace {

TString toString(const Operation& op) {
    switch (op) {
        case Operation::MAPPER:
            return "mapper";
        case Operation::REDUCER:
            return "reducer";
        default:
            throw RuntimeError();
    }
}

namespace option {
    static const TString JOB_COUNT = "job_count";
    static const TString GPU_LIMIT = "gpu_limit";
    static const TString POOL_TREES = "pool_trees";
    static const TString LAYER_PATHS = "layer_paths";
    static const TString MEMORY_LIMIT = "memory_limit";
    static const TString TITLE = "title";
    static const TString RESOURCE_LIMITS = "resource_limits";
    static const TString USER_SLOTS = "user_slots";
    static const TString JOB_IO = "job_io";
    static const TString TABLE_WRITER = "table_writer";
    static const TString MAX_ROW_WEIGHT = "max_row_weight";
    static const TString SCHEDULING_PER_POOL_TREE = "scheduling_options_per_pool_tree";
    static const TString POOL = "pool";
}

static const TString XENIAL_NVIDIA_PORTO_LAYER_PATH
    = "//home/maps/core/mrc/porto_layers/xenial_nvidia-418_cuda-10.1.tar.xz";

static const NYT::TNode GPU_PORTO_LAYERS_PATHS
    = NYT::TNode::CreateList().Add(XENIAL_NVIDIA_PORTO_LAYER_PATH);

} // namespace


YTOpExecutor::MapSpec& YTOpExecutor::MapSpec::AddInput(const NYT::TRichYPath& value) {
    inputs_.push_back(value);
    return *this;
}

YTOpExecutor::MapSpec& YTOpExecutor::MapSpec::AddOutput(const NYT::TRichYPath& value) {
    outputs_.push_back(value);
    return *this;
}

YTOpExecutor::MapSpec::operator NYT::TMapOperationSpec() const {
    NYT::TMapOperationSpec spec;
    for (const NYT::TRichYPath& path : inputs_) {
        spec.AddInput<NYT::TNode>(path);
    }
    for (const NYT::TRichYPath& path : outputs_) {
        spec.AddOutput<NYT::TNode>(path);
    }
    return spec;
}


YTOpExecutor::ReduceSpec& YTOpExecutor::ReduceSpec::AddInput(const NYT::TRichYPath& value) {
    inputs_.push_back(value);
    return *this;
}

YTOpExecutor::ReduceSpec& YTOpExecutor::ReduceSpec::AddOutput(const NYT::TRichYPath& value) {
    outputs_.push_back(value);
    return *this;
}

YTOpExecutor::ReduceSpec& YTOpExecutor::ReduceSpec::ReduceBy(const NYT::TSortColumns& columns) {
    columns_ = columns;
    return *this;
}

YTOpExecutor::ReduceSpec::operator NYT::TReduceOperationSpec() const {
    NYT::TReduceOperationSpec spec;
    for (const NYT::TRichYPath& path : inputs_) {
        spec.AddInput<NYT::TNode>(path);
    }
    for (const NYT::TRichYPath& path : outputs_) {
        spec.AddOutput<NYT::TNode>(path);
    }
    spec.ReduceBy(columns_);
    return spec;
}


YTOpExecutor::Options::Options()
    : runningJobCount_(0)
    , jobCount_(0)
    , memoryLimit_(0)
    , gpuLimit_(0)
    , title_("")
    , maxRowWeight_(0)
{}

YTOpExecutor::Options& YTOpExecutor::Options::RunningJobCount(uint64_t count) {
    runningJobCount_ = count;
    return *this;
}

YTOpExecutor::Options& YTOpExecutor::Options::JobCount(uint64_t count) {
    jobCount_ = count;
    return *this;
}

YTOpExecutor::Options& YTOpExecutor::Options::MemoryLimit(uint64_t limit) {
    memoryLimit_ = limit;
    return *this;
}

YTOpExecutor::Options& YTOpExecutor::Options::UseGPU(uint64_t limit) {
    gpuLimit_ = limit;
    return *this;
}

YTOpExecutor::Options& YTOpExecutor::Options::Title(const TString& title) {
    title_ = title;
    return *this;
}

YTOpExecutor::Options& YTOpExecutor::Options::MaxRowWeight(uint64_t limit) {
    maxRowWeight_ = limit;
    return *this;
}

void YTOpExecutor::Map(
    NYT::IClientBasePtr client,
    const YTOpExecutor::MapSpec& spec,
    ::TIntrusivePtr<NYT::IMapperBase> mapper,
    const YTOpExecutor::Options& options)
{
    client->Map(
        static_cast<NYT::TMapOperationSpec>(spec),
        mapper,
        ToYTOperationOptions(options, Operation::MAPPER)
    );
}

void YTOpExecutor::Reduce(
    NYT::IClientBasePtr client,
    const YTOpExecutor::ReduceSpec& spec,
    ::TIntrusivePtr<NYT::IReducerBase> reducer,
    const YTOpExecutor::Options& options)
{
    client->Reduce(
        static_cast<NYT::TReduceOperationSpec>(spec),
        reducer,
        ToYTOperationOptions(options, Operation::REDUCER)
    );
}


void YTOpExecutor::Initialize(mrc::yt::PoolType poolType, bool useGPU) {
    POOL_TYPE = poolType;
    USE_GPU = useGPU;
}


NYT::TOperationOptions YTOpExecutor::ToYTOperationOptions(
    const YTOpExecutor::Options& options, const Operation& op)
{
    TString opName = toString(op);

    NYT::TNode ytOptions = NYT::TNode::CreateMap();
    if (!options.title_.Empty()) {
        ytOptions[option::TITLE] = options.title_;
    }
    if (options.runningJobCount_ > 0) {
        ytOptions[option::RESOURCE_LIMITS] = NYT::TNode::CreateMap();
        ytOptions[option::RESOURCE_LIMITS][option::USER_SLOTS] = options.runningJobCount_;
    }
    if (options.jobCount_ > 0) {
        ytOptions[option::JOB_COUNT] = options.jobCount_;
    }
    if (options.memoryLimit_ > 0) {
        ytOptions[opName] = NYT::TNode::CreateMap();
        ytOptions[opName][option::MEMORY_LIMIT] = options.memoryLimit_;
    }
    if (options.maxRowWeight_ > 0) {
        ytOptions[option::JOB_IO][option::TABLE_WRITER][option::MAX_ROW_WEIGHT] = options.maxRowWeight_;
    }

    // add porto layer
    if (USE_GPU) {
        if (ytOptions[opName].IsUndefined()) {
            ytOptions[opName] = NYT::TNode::CreateMap();
        }
        ytOptions[opName][option::LAYER_PATHS] = GPU_PORTO_LAYERS_PATHS;
    }

    // add scheduling options
    if (USE_GPU && options.gpuLimit_ > 0) {
        if (ytOptions[opName].IsUndefined()) {
            ytOptions[opName] = NYT::TNode::CreateMap();
        }
        ytOptions[opName][option::LAYER_PATHS] = GPU_PORTO_LAYERS_PATHS;
        ytOptions[opName][option::GPU_LIMIT] = options.gpuLimit_;

        addGpuSchedulingOptions(&ytOptions, POOL_TYPE);
    } else {
        if (options.gpuLimit_ > 0) {
            WARN() << "GPU limit is " << options.gpuLimit_
                   << ", but GPUs are disabled. Operation wiil be peformed on CPU.";
        }

        addCpuSchedulingOptions(&ytOptions, POOL_TYPE);
    }

    if (ytOptions.AsMap().size() > 0) {
        return NYT::TOperationOptions().Spec(ytOptions);
    } else {
        return NYT::TOperationOptions();
    }
}

std::optional<mrc::yt::PoolType> YTOpExecutor::POOL_TYPE = std::nullopt;
bool YTOpExecutor::USE_GPU = false;

} // namespace maps::wiki::autocart::pipeline
