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

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


using namespace NYT;

namespace {

static const TString COLUMN_NAME_DATE       = "date";
static const TString COLUMN_NAME_FEATURE_ID = "feature_id";
static const TString COLUMN_NAME_HEADING    = "heading";
static const TString COLUMN_NAME_IMAGE      = "image";
static const TString COLUMN_NAME_LAT        = "lat";
static const TString COLUMN_NAME_LON        = "lon";
static const TString COLUMN_NAME_OBJECTS    = "objects";

NYT::TNode createSchema() {
    return NYT::TNode::CreateList()
                .Add(NYT::TNode()("name", COLUMN_NAME_DATE)("type", "string"))
                .Add(NYT::TNode()("name", COLUMN_NAME_FEATURE_ID)("type", "int64"))
                .Add(NYT::TNode()("name", COLUMN_NAME_HEADING)("type", "double"))
                .Add(NYT::TNode()("name", COLUMN_NAME_IMAGE)("type", "string"))
                .Add(NYT::TNode()("name", COLUMN_NAME_LAT)("type", "double"))
                .Add(NYT::TNode()("name", COLUMN_NAME_LON)("type", "double"))
                .Add(NYT::TNode()("name", COLUMN_NAME_OBJECTS)("type", "any"));
}

class TDatasetMapper
    : public NYT::IMapper<NYT::TTableReader<NYT::TNode>, NYT::TTableWriter<NYT::TNode>>  {
public:
    TDatasetMapper() = default;

    void Do(NYT::TTableReader<NYT::TNode>* reader, NYT::TTableWriter<NYT::TNode>* writer) override {
        for (; reader->IsValid(); reader->Next()) {
            const NYT::TNode &inpRow = reader->GetRow();

            NYT::TNode outRow;
            outRow  (COLUMN_NAME_FEATURE_ID, inpRow[COLUMN_NAME_FEATURE_ID].AsInt64())
                    (COLUMN_NAME_DATE,       inpRow[COLUMN_NAME_DATE].AsString())
                    (COLUMN_NAME_HEADING,    inpRow[COLUMN_NAME_HEADING].AsDouble())
                    (COLUMN_NAME_IMAGE,      inpRow[COLUMN_NAME_IMAGE].AsString())
                    (COLUMN_NAME_LAT,        inpRow[COLUMN_NAME_LAT].AsDouble())
                    (COLUMN_NAME_LON,        inpRow[COLUMN_NAME_LON].AsDouble())
                    (COLUMN_NAME_OBJECTS,    ConvertObjects(inpRow[COLUMN_NAME_OBJECTS]));
            writer->AddRow(outRow);
        }
    }
private:
    TNode ConvertObjects(const TNode& node) {
        static const TString OBJECT_TYPE_NAME = "type";
        static const TString OBJECT_NUMBER_NAME = "num";

        static const TString PROHIBITORY_MAX_SPEED_PREFIX = "prohibitory_max_speed_";

        const TVector<TNode>& objectList = node.AsList();
        TNode resultList = TNode::CreateList();
        for (const TNode& objectNode : objectList) {
            TNode newNode = objectNode;
            TString type = objectNode[OBJECT_TYPE_NAME].AsString();
            TString number = "";
            if (type.find(PROHIBITORY_MAX_SPEED_PREFIX) == 0) {
                number = type.substr(PROHIBITORY_MAX_SPEED_PREFIX.length());
            }
            newNode(OBJECT_NUMBER_NAME, number);
            resultList.Add(newNode);
        }
        return resultList;
    }
};

REGISTER_MAPPER(TDatasetMapper);

} // namespace

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

    maps::cmdline::Parser parser("Convert traffic signs dataset from one YT table to another, add \'num\' field for traffic sign, auto fill this field for \'prohibitory_max_speed_\' classes");

    maps::cmdline::Option<std::string>
        inputTableName =  parser.string("input")
                                .required()
                                .help("Input YT table name");

    maps::cmdline::Option<std::string>
        outputTableName = parser.string("output")
                                .required()
                                .help("Output YT table name");

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

    INFO() << "Connecting to yt::hahn";
    IClientPtr client = CreateClient("hahn");
    client->Create(outputTableName.c_str(),
                   NYT::NT_TABLE,
                   NYT::TCreateOptions().Attributes(NYT::TNode()("schema", createSchema())));

    client->Map(
        TMapOperationSpec()
            .AddInput<TNode>(inputTableName.c_str())
            .AddOutput<TNode>(outputTableName.c_str()),
        new TDatasetMapper());
    client->Sort(
        TSortOperationSpec()
            .AddInput(outputTableName.c_str())
            .Output(outputTableName.c_str())
            .SortBy({"feature_id"}));

    return EXIT_SUCCESS;
}
catch (const maps::Exception& e) {
    FATAL() << "Worker failed: " << e;
    return EXIT_FAILURE;
}
catch (const std::exception& e) {
    FATAL() << "Worker failed: " << e.what();
    return EXIT_FAILURE;
}
