#include "yt.h"

#include "saas_push.h"

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

#include <saas/library/persqueue/logger/logger.h>

#include <saas/tools/saas_push/key_value.pb.h>

#include <library/cpp/getopt/last_getopt.h>
#include <library/cpp/streams/special/throttle.h>


using namespace NSaasPush;

int YtMain(int argc, const char** argv) {
    NYT::SetLogger(NYT::CreateStdErrLogger(NYT::ILogger::INFO));

    TString configPath;
    TString alias;
    TString ytProxy;
    TString ytTablePath;
    TMaybe<size_t> rateLimit;
    size_t retryLimit;
    size_t threads;

    NLastGetopt::TOpts opts;
    opts.AddHelpOption('h');
    opts.AddLongOption('c', "config", "path to config")
        .Required()
        .RequiredArgument("<path>")
        .StoreResult(&configPath);


    opts.AddLongOption("alias", "Alias")
        .Required()
        .StoreResult(&alias);

    opts.AddLongOption("yt-proxy", "YT proxy")
        .Required()
        .StoreResult(&ytProxy);

    opts.AddLongOption("yt-table-path", "Path to YT table with messages to push")
        .Required()
        .StoreResult(&ytTablePath);

    opts.AddLongOption("rate-limit", "Rate limit per second for sending documents")
        .Handler1T<size_t>([&](size_t rateLimitValue) { rateLimit = rateLimitValue; });

    opts.AddLongOption("retry-limit", "Number of retries for each of the messages")
        .StoreResult(&retryLimit)
        .DefaultValue(10);

    opts.AddLongOption("threads-count", "Number of thread in pool to execute writers")
        .StoreResult(&threads)
        .DefaultValue(1);

    opts.SetFreeArgsMax(0);
    NLastGetopt::TOptsParseResult parseResults(&opts, argc, argv);

    auto config = NSaasPush::ReadConfig(configPath, 0, 0);
    TSearchMapStorage searchMapStorage(
        config.ServerConfig.SearchMapSettings,
        config.ServerConfig.WriterConfig.Services,
        config.ServerConfig.UpdateSearchMapPeriod
    );
    const auto pqLib = MakeAtomicShared<NPersQueue::TPQLib>(config.ServerConfig.PQLibSettings);
    TLog logbrokerLog(config.ServerConfig.LogbrokerLog, config.ServerConfig.LogLevel);
    pqLib->SetLogger(MakeIntrusive<NSaas::TPersQueueLogger<NPersQueue::ILogger>>(logbrokerLog));
    TLog TvmLog(config.ServerConfig.LogbrokerLog, config.ServerConfig.LogLevel);
    TWriter writer{config.ServerConfig.WriterConfig, pqLib, searchMapStorage, TvmLog};

    INFO_LOG << "Waiting until writer is initialized...\n";
    while (!writer.IsInitialized()) {
        Sleep(TDuration::Seconds(1));
    }

    const auto client = NYT::CreateClient(ytProxy);
    const auto rowCount = client->Get(ytTablePath + "/@row_count").AsInt64();
    const auto reader = client->CreateTableReader<TKeyValueProto>(ytTablePath);

    TMaybe<TThrottle> throttle;

    if (rateLimit) {
        throttle.ConstructInPlace(TThrottle::TOptions{*rateLimit, TDuration::Seconds(1)});
    }

    TAtomic writtenCount = 0;

    auto pool = CreateThreadPool(threads, threads, TThreadPool::TParams().SetBlocking(true).SetCatching(false));
    pool->Start(threads, threads);

    for (const auto& iter: *reader) {
        if (throttle) {
            Y_UNUSED(throttle->GetQuota(1));
        }
        pool->SafeAddFunc([&writtenCount, &writer, &alias, &retryLimit, &rowCount, row=iter.GetRow()]() {
            for (size_t retryNumber = 0; retryNumber != retryLimit; ++retryNumber) {
                try {
                    const auto result = writer.Write(alias, row.GetValue());
                    if (result.Written) {
                        break;
                    }
                    INFO_LOG << "Error writing message: " << result.ToString() << '\n';
                } catch (...) {
                    INFO_LOG << "Error writing message: " << CurrentExceptionMessage() << '\n';
                }
                Sleep(TDuration::Seconds(1));
            }
            AtomicIncrement(writtenCount);
            if (int written = AtomicGet(writtenCount); written % 1000 == 0) {
                INFO_LOG << "Written " << written << " out of " << rowCount << " (" << (100.0 * written / rowCount) << "%)\n";
            }
        });
    }

    INFO_LOG << "Done\n";

    return 0;
}
