#include <maps/libs/log8/include/log8.h>
#include <maps/libs/common/include/exception.h>
#include <maps/libs/cmdline/include/cmdline.h>

#include <maps/libs/json/include/builder.h>

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

#include <random>
#include <vector>
#include <fstream>

using namespace maps;

namespace {

std::string getFeatureURL(int64_t featureId) {
    return "https://core-nmaps-mrc-browser.maps.yandex.ru/feature/" + std::to_string(featureId) + "/image";
}

class Task {
public:
    Task(const NYT::TNode& node)
        : featureId_(node["feature_id"].AsInt64())
        , featureURL_(getFeatureURL(featureId_))
    {
    }

    void toJson(json::ObjectBuilder& builder) const {
        builder["inputValues"] = [&](json::ObjectBuilder b) {
            b["feature_id"] = featureId_;
            b["image"] = featureURL_;
        };
    }

private:
    int64_t featureId_;
    std::string featureURL_;
};

} // namespace

int main(int argc, const char** argv)
try {
    NYT::Initialize(argc, argv);

    cmdline::Parser parser("Prepare toloka tasks in JSON");

    cmdline::Option<std::string> ytPath = parser.string("yt_path")
        .required()
        .help("Path to YT table with features");

    cmdline::Option<std::string> tasksPath = parser.string("tasks_path")
        .required()
        .help("Path to output json with tasks");

    cmdline::Option<size_t> tasksCount = parser.size_t("count")
        .required()
        .help("Tasks count in output json");

    cmdline::Option<bool> doShuffle = parser.flag("random_shuffle")
        .help("Shuffle rows in YT tables");

    cmdline::Option<size_t> randomSeed = parser.size_t("seed")
        .defaultValue(42)
        .help("Seed for random number generator");

    parser.parse(argc, const_cast<char**>(argv));

    INFO() << "Create YT client: yt::hahn";
    NYT::IClientPtr client = NYT::CreateClient("hahn");

    std::vector<Task> tasks;

    INFO() << "Reading tasks from YT table: " << ytPath;
    NYT::TTableReaderPtr<NYT::TNode> reader
        = client->CreateTableReader<NYT::TNode>(NYT::TYPath(ytPath));
    if (!doShuffle) {
        INFO() << "Reading " << tasksCount << " rows from YT table";
        size_t addedCount = 0;
        for (; addedCount < tasksCount && reader->IsValid(); reader->Next()) {
            try {
                Task task(reader->GetRow());
                tasks.push_back(task);
                addedCount++;
            } catch (const maps::RuntimeError& e) {
                INFO() << "Failed to create task";
            }
        }
    } else {
        INFO() << "Reading all rows from YT table";
        std::vector<NYT::TNode> rows;
        for (; reader->IsValid(); reader->Next()) {
            rows.push_back(reader->GetRow());
        }
        INFO() << "Shuffling rows";
        std::mt19937 rndGen{randomSeed};
        INFO() << "Reading " << tasksCount << " rows";
        std::shuffle(rows.begin(), rows.end(), rndGen);
        size_t addedCount = 0;
        for (auto it = rows.begin(); addedCount < tasksCount && it != rows.end() ; it++) {
            try {
                Task task(*it);
                tasks.push_back(task);
                addedCount++;
            } catch (const maps::Exception& e) {
                INFO() << "Failed to create task: " << e.what();
            }
        }
    }

    INFO() << "Dumping tasks to json: " << tasksPath;
    std::ofstream ofs(tasksPath);
    json::Builder builder(ofs);
    builder << [&](json::ArrayBuilder b) {
        for (const Task& task : tasks) {
            b << [&](json::ObjectBuilder b) {
                task.toJson(b);
            };
        }
    };
    ofs.close();

    return EXIT_SUCCESS;
}
catch (const maps::Exception& e) {
    INFO() << e;
    return EXIT_FAILURE;
}
catch (const std::exception& e) {
    INFO() << e.what();
    return EXIT_FAILURE;
}
catch (...) {
    INFO() << "Caught unknown exception";
    return EXIT_FAILURE;
}
