#include <saas/api/indexing_client/client.h>
#include <saas/api/mr_client/client.h>
#include <saas/api/indexing_client/sloth.h>
#include <saas/util/logging/tskv_log.h>
#include <saas/util/queue.h>

#include <mapreduce/lib/init.h>

#include <library/cpp/getopt/last_getopt.h>
#include <library/cpp/logger/global/global.h>
#include <library/cpp/json/json_reader.h>

#include <util/stream/file.h>
#include <library/cpp/deprecated/atomic/atomic_ops.h>
#include <util/random/random.h>

class IBuilderAgent {
public:
    virtual ~IBuilderAgent() {}
    using TPtr = TAtomicSharedPtr<IBuilderAgent>;

    virtual TAtomicSharedPtr<NSaas::TAction> Build() const = 0;
};

class TFileSourceBuilder: public IBuilderAgent {
private:
    const TString FileName;
    mutable TFileInput BFI;
    TMutex Mutex;
public:

    TFileSourceBuilder(const TString& fileName)
        : FileName(fileName)
        , BFI(FileName)
    {

    }

    TAtomicSharedPtr<NSaas::TAction> Build() const override {
        TString line;
        {
            TGuard<TMutex> g(Mutex);
            line = BFI.ReadLine();
            if (!line) {
                return nullptr;
            }
        }
        TMap<TString, TString> parseResult;
        NUtil::TTSKVRecordParser::Parse<'\t', '='>(line, parseResult);
        NSaas::TAction action;
        NJson::TJsonValue json;
        TStringInput is(parseResult["message"]);
        CHECK_WITH_LOG(NJson::ReadJsonTree(&is, true, &json));

        auto result = MakeAtomicShared<NSaas::TAction>();
        result->ParseFromJson(json);
        return result;
    }
};

class TLaasBuilder: public IBuilderAgent {
private:
    ui64 Count;
    mutable TAtomic CurrentCounter = 0;
public:
    TLaasBuilder(const ui64 count)
        : Count(count)
    {
    }

    TAtomicSharedPtr<NSaas::TAction> Build() const override {

        ui64 current = AtomicIncrement(CurrentCounter);
        if (current > Count)
            return nullptr;

        ui16 i1 = RandomNumber<ui16>();
        ui16 i2 = RandomNumber<ui16>();
        ui16 i3 = RandomNumber<ui16>();
        ui16 i4 = RandomNumber<ui16>();

        TAtomicSharedPtr<NSaas::TAction> action = new NSaas::TAction;

        action->SetActionType(NSaas::TAction::atModify).GetDocument().SetUrl(ToString(i1) + "." + ToString(i2) + "." + ToString(i3) + "." + ToString(i4)).SetBody("abc").SetRealtime(true);

        return action;
    }
};

struct TSenderAgentOption {
    ui32 ReindexAttemptions = 1;
    TDuration MinSleepDuration = TDuration::MilliSeconds(10);
    TDuration MaxSleepDuration = TDuration::Seconds(10);
};

class TSenderAgent: public IObjectInQueue {
private:
    IBuilderAgent::TPtr Builder;
    NSaas::TIndexingClient& Client;
    TSenderAgentOption Option;
public:

    TSenderAgent(NSaas::TIndexingClient& client, IBuilderAgent::TPtr builder, const TSenderAgentOption &option)
        : Client(client)
        , Option(option)
    {
        Builder = builder;
    }

    void Process(void* /*ThreadSpecificResource*/) override {
        TAtomicSharedPtr<NSaas::TAction> action;
        // We start with a zero sleep between requests, then increase it on each failure and decrease on success.
        // The requests-per-seconds rate is expected to converge around the indexer proxy throughput.
        TSlothExp sloth(TDuration::Zero(), Option.MinSleepDuration, Option.MaxSleepDuration);
        while (!!(action = Builder->Build())) {
            for (ui32 i = 0; i < 1 + Option.ReindexAttemptions; ++i) {  // one initial attempt + number of retries
                TDuration sleepedFor = sloth.Sleep();
                DEBUG_LOG << "sleeped for " << sleepedFor.MilliSeconds() << "ms" << Endl;
                NSaas::TSendResult sendResult = Client.Send(*action);
                DEBUG_LOG << sendResult.GetMessage() << Endl;
                if (sendResult.GetCode() == NSaas::TSendResult::srOK) {
                    sloth.DecDuration();
                    break;
                }
                sloth.IncDuration();
            }
        }
    }
};

struct TMRContext {
    TString MRServer;
    TString SrcTable;
    TString DstTable;
    int Jobs = 200;

    bool IsValid() {
        return !SrcTable.empty() && !DstTable.empty();
    }
};

void RunYTClient(const NSaas::TSaasIndexerOption& options, const TMRContext& mrCtx) {
    auto client = NYT::CreateClient(mrCtx.MRServer);
    client->Create(mrCtx.DstTable, NYT::NT_TABLE, NYT::TCreateOptions().IgnoreExisting(true));
    NYT::TMapOperationSpec spec;
    spec.AddInput<NYT::TNode>(mrCtx.SrcTable);
    spec.AddOutput<NYT::TNode>(mrCtx.DstTable);
    client->Map(spec, new NSaas::TSaasYTMapper(options));
}

void RunMRClient(const NSaas::TSaasIndexerOption& options, const TMRContext& mrCtx) {
    NMR::NLog::SetStderrLevel(5);

    NMR::TServer server(mrCtx.MRServer);
    NMR::TMRParams params;
    params.SetJobCount(mrCtx.Jobs);
    params.AddInputTable(mrCtx.SrcTable);
    params.AddOutputTable(mrCtx.DstTable);
    server.Map(params, new NSaas::TSaasIndexerMap(options));

    NMR::TClient client(server);
    NMR::TTable table(client, mrCtx.DstTable);

    ui32 i = 0;
    for (NMR::TTableIterator it = table.Begin(); it.IsValid(); ++it) {
        TString key = it.GetKey().AsString();
        TString value = it.GetValue().AsString();
        Cout << key << '\t' << value << Endl;
        if (i == 10)
            break;
        ++i;
    }
}

void RunLocal(const NSaas::TSaasIndexerOption& ctx, ui32 count, const TString& fileName) {
    TRTYMtpQueue queue;
    IBuilderAgent::TPtr builderAgent;
    if (!fileName) {
        builderAgent = new TLaasBuilder(count);
    } else {
        builderAgent = new TFileSourceBuilder(fileName);
    }
    queue.Start(ctx.Threads);
    NSaas::TIndexingClient client(ctx.Host, ctx.Port, ctx.GetIndexUrl());

    TSenderAgentOption senderOpt { ctx.ReindexAttemptions };
    for (ui32 i = 0; i < ctx.Threads; ++i) {
        queue.SafeAddAndOwn(THolder(new TSenderAgent(client, builderAgent, senderOpt)));
    }
    queue.Stop();
}

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

    NLastGetopt::TOpts opts;
    opts.AddHelpOption();

    NSaas::TSaasIndexerOption ctx;
    TString fileName;
    opts.AddCharOption('h', "host").StoreResult(&ctx.Host).DefaultValue("saas-indexerproxy-prestable.yandex.net");
    opts.AddCharOption('p', "port").StoreResult(&ctx.Port).DefaultValue("80");
    opts.AddCharOption('k', "indexkey").StoreResult(&ctx.IndexKey).Required();
    opts.AddCharOption('f', "ammo filename").StoreResult(&fileName).DefaultValue("");
    opts.AddCharOption('l', "log level").StoreResult(&ctx.LogLevel).DefaultValue("6");
    opts.AddCharOption('r', "number of reindex attemptions").StoreResult(&ctx.ReindexAttemptions).DefaultValue("1");
    opts.AddLongOption("prefix").StoreResult(&ctx.ProcessorOptions.Prefix);
    opts.AddLongOption("action").StoreResult(&ctx.ProcessorOptions.ActionType);

    TString runMode;
    opts.AddLongOption("mode").StoreResult(&runMode).DefaultValue("local");

    ui32 count = 0;
    opts.AddCharOption('c', "messages count").StoreResult(&count);
    opts.AddCharOption('t', "subqueue threads count").StoreResult(&ctx.Threads);
    opts.AddLongOption("queue-size", "subqueue max size").StoreResult(&ctx.QueueSize);
    opts.AddLongOption("compression").StoreResult(&ctx.ProcessorOptions.Compression).DefaultValue("none");
    opts.AddLongOption("rty_adapter").StoreResult(&ctx.RtyIndexAdapter).DefaultValue("json");
    opts.AddLongOption("timestamp").StoreResult(&ctx.ProcessorOptions.Timestamp);

    TMRContext mrContext;
    opts.AddLongOption("mr-server").StoreResult(&mrContext.MRServer).DefaultValue("plato.yt.yandex.net:80");
    opts.AddLongOption("mr-input").StoreResult(&mrContext.SrcTable);
    opts.AddLongOption("mr-output").StoreResult(&mrContext.DstTable);
    opts.AddLongOption("mr-processor").StoreResult(&ctx.Processor);
    opts.AddLongOption("mr-processor-env").StoreResult(&ctx.ProcessorOptions.OtherOptions);
    opts.AddLongOption("mr-jobs").StoreResult(&mrContext.Jobs);

    NLastGetopt::TOptsParseResult parseOpt(&opts, argc, argv);
    DoInitGlobalLog("console", ctx.LogLevel, false, false);

    if (!parseOpt.Has("timestamp")) {
        ctx.ProcessorOptions.Timestamp = TInstant::Now().Seconds();
    }

    if (runMode == "local") {
        RunLocal(ctx, count, fileName);
    } else if (runMode == "mr" || runMode == "yt") {
        if (!mrContext.IsValid())
            ythrow yexception() << "invalid mr parameters";
        if (runMode == "yt") {
            RunYTClient(ctx, mrContext);
        } else {
            RunMRClient(ctx, mrContext);
        }
    } else {
        ythrow yexception() << "Unknown rum mode";
    }

    return 0;
}
