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

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

#include <util/generic/size_literals.h>

#include <mapreduce/yt/util/temp_table.h>
#include <mapreduce/yt/interface/client.h>
#include <mapreduce/yt/util/ypath_join.h>

#include <fstream>

namespace {

static const char* FT_ID = "ft_id";
static const char* SHAPE = "shape";
static const char* FT_TYPE_ID = "ft_type_id";

class FTTypeIdFilter
    : public NYT::IMapper<NYT::TTableReader<NYT::TNode>,
                          NYT::TTableWriter<NYT::TNode>>
{
public:
    FTTypeIdFilter() = default;
    FTTypeIdFilter(int ftTypeId)
        : ftTypeId_(ftTypeId)
    {
    }

    Y_SAVELOAD_JOB(ftTypeId_);

    void Do(NYT::TTableReader<NYT::TNode>* reader,
            NYT::TTableWriter<NYT::TNode>* writer) override
    {
        for (; reader->IsValid(); reader->Next()) {
            const NYT::TNode& row = reader->GetRow();
            if (row[FT_TYPE_ID].AsInt64() == ftTypeId_) {
                writer->AddRow(row);
            }
        }
    }

private:
    int ftTypeId_;
};

REGISTER_MAPPER(FTTypeIdFilter);


class MergerByFTId
    : public NYT::IReducer<NYT::TTableReader<NYT::TNode>,
                           NYT::TTableWriter<NYT::TNode>>
{
public:
    void Do(NYT::TTableReader<NYT::TNode>* reader,
            NYT::TTableWriter<NYT::TNode>* writer) override
    {
        static const size_t FT_TYPE_ID_INDEX = 0;

        std::optional<NYT::TNode> ftTypeNode;
        std::optional<NYT::TNode> ftGeomNode;
        for (; reader->IsValid(); reader->Next()) {
            const NYT::TNode& row = reader->GetRow();
            size_t index = reader->GetTableIndex();
            if (index == FT_TYPE_ID_INDEX) {
                ftTypeNode = row;
            } else {
                ftGeomNode = row;
            }
        }
        if (ftTypeNode && ftGeomNode) {
            NYT::TNode node = NYT::TNode::CreateMap();
            node[FT_ID] = (*ftTypeNode)[FT_ID];
            node[FT_TYPE_ID] = (*ftTypeNode)[FT_TYPE_ID];
            node[SHAPE] = (*ftGeomNode)[SHAPE];
            writer->AddRow(node);
        }
    }
};

REGISTER_REDUCER(MergerByFTId);


NYT::TOperationOptions setMaxRowWeight(size_t maxWeight) {
    return NYT::TOperationOptions()
        .Spec(NYT::TNode::CreateMap()
            ("job_io", NYT::TNode::CreateMap()
                ("table_writer", NYT::TNode::CreateMap()
                    ("max_row_weight", maxWeight)
                )
            )
        );
}

void filterByFTTypeId(
    NYT::IClientPtr client,
    const TString& ymapsdfYTFolder,
    const TString& outputYTPath,
    int ftTypeId)
{
    NYT::TNode::TListType regionNodes = client->List(ymapsdfYTFolder);
    NYT::TMapOperationSpec spec;
    for (int i = 0; i < regionNodes.ysize(); i++) {
        TString region = regionNodes[i].AsString();
        spec.AddInput<NYT::TNode>(NYT::JoinYPaths(ymapsdfYTFolder, region, "ft"));
    }
    spec.AddOutput<NYT::TNode>(outputYTPath);

    client->Map(
        spec,
        new FTTypeIdFilter(ftTypeId)
    );
}

void mergeByFTId(
    NYT::IClientPtr client,
    const TString& ftTypeIdYTPath,
    const TString& ymapsdfYTFolder,
    const TString& outputYTPath)
{
    NYT::TNode::TListType regionNodes = client->List(ymapsdfYTFolder);
    NYT::TReduceOperationSpec spec;
    spec.AddInput<NYT::TNode>(ftTypeIdYTPath);
    for (int i = 0; i < regionNodes.ysize(); i++) {
        TString region = regionNodes[i].AsString();
        spec.AddInput<NYT::TNode>(NYT::JoinYPaths(ymapsdfYTFolder, region, "ft_geom"));
    }
    spec.AddOutput<NYT::TNode>(outputYTPath);
    spec.ReduceBy({FT_ID});

    client->Reduce(
        spec,
        new MergerByFTId(),
        setMaxRowWeight(128_MB)
    );
}

} // namespace

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

    maps::cmdline::Parser parser("Make dataset by ft_type_id");

    maps::cmdline::Option<std::string> ymapsdfYTFolder = parser.string("ymapsdf_folder")
        .required()
        .help("Path to folder with YMapsDF tables");

    maps::cmdline::Option<std::size_t> ftTypeId = parser.size_t("ft_type_id")
        .required()
        .help("Selected ft_type_id");

    maps::cmdline::Option<std::string> outputJsonPath = parser.string("output")
        .required()
        .help("Path to output json file with dataset");

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

    INFO() << "Connecting to yt::hahn";
    NYT::IClientPtr ytClient = NYT::CreateClient("hahn");

    INFO() << "Filtering by ft type id: " << ftTypeId;
    NYT::TTempTable ftTypeIdTable(ytClient);
    filterByFTTypeId(
        ytClient,
        TString(ymapsdfYTFolder),
        ftTypeIdTable.Name(),
        ftTypeId
    );

    ytClient->Sort(
        NYT::TSortOperationSpec()
            .AddInput(ftTypeIdTable.Name())
            .Output(ftTypeIdTable.Name())
            .SortBy({FT_ID})
    );

    INFO() << "Merging tables by " << FT_ID;
    NYT::TTempTable resultTable(ytClient);
    mergeByFTId(
        ytClient,
        ftTypeIdTable.Name(),
        TString(ymapsdfYTFolder),
        resultTable.Name()
    );

    INFO() << "Saving dataset into json: " << outputJsonPath;
    NYT::TTableReaderPtr<NYT::TNode> reader
        = ytClient->CreateTableReader<NYT::TNode>(resultTable.Name());
    std::ofstream ofs(outputJsonPath);
    REQUIRE(ofs.is_open(), "Failed to create file: " + outputJsonPath);
    maps::json::Builder builder(ofs);
    builder << [&](maps::json::ArrayBuilder b) {
        for (; reader->IsValid(); reader->Next()) {
            const NYT::TNode& row = reader->GetRow();
            b << [&](maps::json::ObjectBuilder b) {
                b[FT_ID] = row[FT_ID].AsInt64();
                b[FT_TYPE_ID] = row[FT_TYPE_ID].AsInt64();
                b[SHAPE] = row[SHAPE].AsString();
            };
        }
    };
    ofs.close();

    INFO() << "Done!";

    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;
}
