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

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

#include <fstream>

namespace {

static const TString SHAPE = "shape";
static const TString STATE = "state";

static const std::string STATE_YES = "yes";
static const std::string STATE_NO = "no";

using HexWKB = std::string;

struct Item {
    Item(const HexWKB& hexWKB, const std::string& state)
        : hexWKB(hexWKB),
          state(state) {}

    HexWKB hexWKB;
    std::string state;
};

using Dataset = std::vector<Item>;

void loadDataFromTable(const NYT::ITransactionPtr txn,
                       const TString& tableYTPath,
                       Dataset* dataset) {
    auto reader = txn->CreateTableReader<NYT::TNode>(tableYTPath);
    for (; reader->IsValid(); reader->Next()) {
        const auto& row = reader->GetRow();
        HexWKB hexWKB = row[SHAPE].AsString();
        std::string state = row[STATE].AsString();
        if (STATE_YES == state || STATE_NO == state) {
            dataset->emplace_back(hexWKB, state);
        }
    }
}

void saveDataset(const std::string& datasetPath,
                 const Dataset& dataset) {
    std::ofstream ofs(datasetPath);
    for (const Item& item : dataset) {
        ofs << item.hexWKB << "," << item.state << "\n";
    }
    ofs.close();
}

} // namespace

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

    maps::cmdline::Parser parser("Downloading dataset for AutoToloker");

    auto tableYTPath = parser.string("table_path")
        .required()
        .help("Path to table in YT");

    auto datasetPath = parser.string("dataset_path")
        .required()
        .help("Path to dataset");

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

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

    NYT::ITransactionPtr txn = client->StartTransaction();

    INFO() << "Loading dataset from table: " << tableYTPath;
    Dataset dataset;
    loadDataFromTable(txn, TString(tableYTPath), &dataset);

    txn->Commit();

    INFO() << "Saving dataset: " << datasetPath;
    saveDataset(datasetPath, dataset);

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