#include <util/generic/size_literals.h>
#include <mapreduce/yt/interface/client.h>
#include <mapreduce/yt/common/config.h>
#include <datacloud/dev_utils/data/custom_data_utils.h>
#include <datacloud/features/dssm/sample_app/data.pb.h>

using namespace NYT;

namespace Datacloud {
namespace Retro {

class TGetRetroLogsReducer : public IReducer<
    TTableReader<::google::protobuf::Message>,
    TTableWriter<TOutputLogRecord> > 
{
public:
    void Do(TReader* reader, TWriter* writer) override
    {
        i64 maxTimestamp = 0;
        TVector<TString> extIds;
        for (; reader->IsValid(); reader->Next()) {
            const auto tableIndex = reader->GetTableIndex();
            if (tableIndex == 0) {
                const auto& row = reader->GetRow<TInputYuidRecord>();
                maxTimestamp = row.GetTimestamp();
                extIds.emplace_back(row.GetExternalId());
            } else if (!extIds.empty()) {
                const auto& row = reader->GetRow<TInputLogRecord>();
                const auto timestamp = row.GetTimestamp();
                if (timestamp > maxTimestamp) {
                    continue;
                }
                TOutputLogRecord outputRecord;
                outputRecord.SetYuid(row.GetYuid());
                outputRecord.SetTimestamp(timestamp);
                outputRecord.SetTitle(row.GetTitle());
                outputRecord.SetUrl(row.GetUrl());
                for (const auto& extId : extIds) {
                    outputRecord.SetExternalId(extId);
                    writer->AddRow(outputRecord);
                }
            } else {
                break;
            }
        }
    }
};
REGISTER_REDUCER(TGetRetroLogsReducer);

void GetRetroLogs(
        const TString& ytToken,
        const TString& cluster,
        const TString& inputYuidTable,
        const TVector<TString>& inputLogTables,
        const TString& outputTable) {
    auto client = NYT::CreateClient(
        cluster, NYT::TCreateClientOptions().Token(ytToken));

    const auto outputSchema = TTableSchema()
        .AddColumn(TColumnSchema().Type(VT_STRING).Name("external_id"))
        .AddColumn(TColumnSchema().Type(VT_STRING).Name("yuid"))
        .AddColumn(TColumnSchema().Type(VT_STRING).Name("title"))
        .AddColumn(TColumnSchema().Type(VT_STRING).Name("url"))
        .AddColumn(TColumnSchema().Type(VT_INT64).Name("timestamp"));

    auto spec = TJoinReduceOperationSpec()
        .JoinBy({"yuid"})
        .AddInput<TInputYuidRecord>(
            TRichYPath(inputYuidTable)
            .Foreign(true));
    for (const auto& table : inputLogTables) {
        spec.AddInput<TInputLogRecord>(table);
    }
    spec.AddOutput<TOutputLogRecord>(TRichYPath(outputTable)
        .Schema(outputSchema)
        .OptimizeFor(EOptimizeForAttr::OF_SCAN_ATTR));
    client->JoinReduce(
        spec,
        new Datacloud::Retro::TGetRetroLogsReducer
    );
    client->Sort(
       TSortOperationSpec()
           .AddInput(outputTable)
           .Output(outputTable)
           .SortBy({"external_id", "yuid"}));
}

}  // namespace Retro
}  // namespace Datacloud

// 
// int main(int argc, const char** argv) {
//     Initialize(argc, argv);
//     auto client = CreateClient("hahn");
// 
//     const TString inputYuidTable = "//projects/scoring/beeline/XPROD-1417/input_yuid";
//     const TVector<TString> inputLogTables = {
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-01",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-02",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-03",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-04",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-05",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-06",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-07",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-08",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-09",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-10",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-11",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-12",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-13",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-14",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-15",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-16",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-17",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-18",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-19",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-20",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-21",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-22",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-23",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-24",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-25",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-26",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-27",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-28",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-29",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-30",
//         "//home/x-products/production/datacloud/grep/spy_log/2016-07-31",
//     };
//     const TString outputTable = "//tmp/r9-sample-retrp-grep";
// 
//     const auto outputSchema = TTableSchema()
//         .AddColumn(TColumnSchema().Type(VT_STRING).Name("external_id"))
//         .AddColumn(TColumnSchema().Type(VT_STRING).Name("yuid"))
//         .AddColumn(TColumnSchema().Type(VT_INT64).Name("timestamp"))
//         .AddColumn(TColumnSchema().Type(VT_STRING).Name("title"))
//         .AddColumn(TColumnSchema().Type(VT_STRING).Name("url"));
//     const auto sortedTable = "//tmp/r9-sorted-ext-table";
// 
//     // client->Sort(
//     //    TSortOperationSpec()
//     //        .AddInput(inputYuidTable)
//     //        .Output(sortedTable)
//     //       .SortBy({"yuid"}));
// 
//     auto spec = TJoinReduceOperationSpec()
//         .JoinBy({"yuid"})
//         .AddInput<TInputYuidRecord>(
//             TRichYPath(sortedTable)
//             .Foreign(true));
//     for (const auto& table : inputLogTables) {
//         spec.AddInput<TInputLogRecord>(table);
//     }
//     spec.AddOutput<TOutputLogRecord>(TRichYPath(outputTable)
//         .Schema(outputSchema)
//         .OptimizeFor(EOptimizeForAttr::OF_SCAN_ATTR));
// 
//     client->JoinReduce(
//         spec,
//         new Datacloud::Retro::TGetRetroLogsReducer
//     );
// 
//     client->Sort(
//        TSortOperationSpec()
//            .AddInput(outputTable)
//            .Output(outputTable)
//            .SortBy({"external_id", "yuid"}));
// 
//     return 0;
// }
