#include <balancer/production/x/collect_cookie_stats/lib/cookie.pb.h>
#include <balancer/production/x/collect_cookie_stats/lib/cookie_extractor.h>

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

#include <library/cpp/getopt/last_getopt.h>
#include <library/cpp/yson/node/node_io.h>

#include <util/string/split.h>
#include <util/string/join.h>

#include <contrib/libs/re2/re2/re2.h>

namespace NCookie {
    const ui32 CookieStats = 0;
    const ui32 CookieHeaders = 1;
    const ui32 SetCookies = 2;

    class TCookieMap : public NYT::IMapper<NYT::TTableReader<NYT::TNode>, NYT::TTableWriter<NYT::Message>> {
    public:
        void Do(TReader* reader, TWriter* writer) override {
            THashMap<TString, TCookie> cookies;

            for (auto& cursor : *reader) {
                THashMap<TString, TCookie> localCookies;
                ui32 cookieCount = 0;
                ui32 cookieSize = 0;
                const auto& row = cursor.GetRow();
                auto hostPath = Extractor_.Apply(row, [&](const NCookie::TCookie& c) {
                    cookieCount += 1;
                    cookieSize += c.GetName().size() + 3 + c.GetValueLenMax();
                    MergeCookies(localCookies[c.GetName()], c, true);
                }, [&](const NCookie::TSetCookie& sc) {
                    writer->AddRow<TSetCookie>(sc, SetCookies);
                });

                for (auto&& c : localCookies) {
                    MergeCookies(cookies[c.first], c.second, false);
                }

                if (cookieCount) {
                    TCookieHeader header;
                    header.SetReqHost(hostPath.Host);
                    header.SetReqPath(hostPath.Path);
                    header.SetCount(cookieCount);
                    header.SetSize(cookieSize);
                    writer->AddRow<TCookieHeader>(header, CookieHeaders);
                }
            }

            IterateCookies(cookies, [&](const TCookie& c) {
                writer->AddRow<TCookie>(c, CookieStats);
            });
        }

    private:
        TSetCookieExtractor Extractor_;
    };
    REGISTER_MAPPER(TCookieMap)


    class TCookieReduce : public NYT::IReducer<NYT::TTableReader<TCookie>, NYT::TTableWriter<TCookie>> {
    public:
        void Do(TReader* reader, TWriter* writer) override {
            THashMap<TString, TCookie> cookies;

            for (auto& cursor : *reader) {
                const auto& row = cursor.GetRow();
                MergeCookies(cookies[row.GetName()], row, false);
            }

            IterateCookies(cookies, [&](const TCookie& c) {
                writer->AddRow(c);
            });
        }
    };
    REGISTER_REDUCER(TCookieReduce);
}


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

    TString cluster;
    TString outPrefix;

    using namespace NLastGetopt;
    TOpts opts = NLastGetopt::TOpts::Default();
    opts.AddLongOption('c', "cluster").StoreResult(&cluster).RequiredArgument("CLUSTER").Required();
    opts.AddLongOption('o', "out-prefix").StoreResult(&outPrefix).RequiredArgument("TABLE_PREFIX").Required();

    opts.SetFreeArgsMin(1);
    opts.SetFreeArgTitle(0, "IN_TABLE", " ");
    opts.SetFreeArgDefaultTitle("IN_TABLES", "...");

    TOptsParseResult res(&opts, argc, argv);

    TVector<TString> inTables = res.GetFreeArgs();
    const auto cookieStats = outPrefix + "/cookie_stats";
    const auto cookieHeaders = outPrefix + "/cookie_headers";
    const auto setCookies = outPrefix + "/set_cookies";

    auto client = NYT::CreateClient(cluster);

    auto mapSpec = NYT::TMapOperationSpec();
    for (auto&& table : inTables) {
        mapSpec.AddInput<NYT::TNode>(table);
    }
    mapSpec.SetOutput<NCookie::TCookie>(NCookie::CookieStats, NYT::WithSchema<NCookie::TCookie>(cookieStats));
    mapSpec.SetOutput<NCookie::TCookieHeader>(NCookie::CookieHeaders, NYT::WithSchema<NCookie::TCookieHeader>(cookieHeaders));
    mapSpec.SetOutput<NCookie::TSetCookie>(NCookie::SetCookies, NYT::WithSchema<NCookie::TSetCookie>(setCookies));

    Cout << "map " << JoinSeq(",", inTables) << " -> " << " " << JoinSeq(",", {cookieStats, cookieHeaders, setCookies}) << Endl;
    client->Map(mapSpec, MakeIntrusive<NCookie::TCookieMap>());
    Cout << "done" << Endl;

    Cout << "sort " << cookieStats << " -> " << cookieStats << Endl;
    client->Sort(cookieStats, cookieStats, "name");
    Cout << "done" << Endl;

    auto reduceSpec = NYT::TReduceOperationSpec();
    reduceSpec.ReduceBy("name");
    reduceSpec.AddInput<NCookie::TCookie>(cookieStats);
    reduceSpec.AddOutput<NCookie::TCookie>(NYT::WithSchema<NCookie::TCookie>(cookieStats, "name"));

    Cout << "reduce " << cookieStats << " -> " << cookieStats << Endl;
    client->Reduce(reduceSpec, MakeIntrusive<NCookie::TCookieReduce>());
    Cout << "done" << Endl;
}
