#include "client.h"
#include "options.h"

#include <saas/api/action.h>
#include <saas/protos/reply.pb.h>
#include <saas/protos/rtyserver.pb.h>
#include <saas/library/indexer_protocol/sender_neh.h>
#include <saas/library/indexer_protocol/protocol.h>

#include <library/cpp/balloc/optional/operators.h>
#include <library/cpp/logger/global/global.h>

namespace {
    class TQueuedAction : public IObjectInQueue {
    public:
        TQueuedAction(
            NFusion::TActionPtr action,
            std::function<void(const NRTYServer::TReply&)>&& handler,
            const NFusion::TBaseIndexingClient& client
            )
            : Action(action)
            , Handler(handler)
            , Client(client)
        {
        }

        void Process(void*) override {
            THolder<TQueuedAction> cleanup(this);
            Handler(Client.Send(Action));
        }

    private:
        NFusion::TActionPtr Action;
        std::function<void(const NRTYServer::TReply&)> Handler;
        const NFusion::TBaseIndexingClient& Client;
    };

    bool ShouldRetry(const NRTYServer::TReply& reply) {
        switch (reply.GetStatus()) {
        case NRTYServer::TReply::NOTNOW:
        case NRTYServer::TReply::DATA_ACCEPTED:
            return true;
        default:
            return false;
        }
    }
}

NFusion::TBaseIndexingClient::TBaseIndexingClient(const TString& scheme, const TString& host, ui16 port, const TIndexingClientOptions& options)
    : Options(options)
    , Queue(this, "IndexClient")
    , NehSender(MakeHolder<NRTYServer::TNehSender>(scheme, host, port))
    , Quota(options.AverageRate, options.BurstRate, nullptr, nullptr, nullptr)
    , LastMessageId(0)
    , InFlight(0)
    , Active(true)
{
    Queue.Start(options.Threads, options.BurstRate * Max<size_t>(options.InterAttemptTimeout.Seconds(), 1));
}

NFusion::TBaseIndexingClient::~TBaseIndexingClient() {
    Shutdown();
}

void NFusion::TBaseIndexingClient::Send(TActionPtr action, std::function<void(const NRTYServer::TReply&)> handler, bool useQuota) const {
    CHECK_WITH_LOG(NehSender);
    if (!action) {
        throw yexception() << "Null action has been passed";
    }

    if (!action->GetId())
        action->SetId(AtomicIncrement(LastMessageId));

    if (useQuota) {
        Quota.Use(1, true);
    }

    auto d = MakeHolder<TQueuedAction>(action, std::move(handler), *this);
    while (!Queue.Add(d.Get())) {
        Sleep(Options.InterAttemptTimeout);
    }
    Y_UNUSED(d.Release());
}

NRTYServer::TReply NFusion::TBaseIndexingClient::Send(TActionPtr action) const {
    NRTYServer::TReply reply;
    reply.SetStatus(NRTYServer::TReply::NOTNOW);

    if (!action->GetId())
        action->SetId(AtomicIncrement(LastMessageId));

    AtomicIncrement(InFlight);
    try {
        ui32 attempt = 0;
        while (Active && attempt < Options.Attempts) {
            CHECK_WITH_LOG(action);
            auto context = NehSender->Send(action->ToProtobuf());
            while (context && Active && attempt++ < Options.Attempts && !NehSender->Recv(context, reply, Options.Timeout)) {
                DEBUG_LOG << "Indexing " << action->GetActionDescription() << " timeouted" << Endl;
            }

            if (!ShouldRetry(reply)) {
                break;
            }

            if (Options.InterAttemptTimeout != TDuration::Zero()) {
                Sleep(Options.InterAttemptTimeout);
            }
            DEBUG_LOG << "Retrying to index " << action->GetActionDescription() << Endl;
        }
        if (!Active) {
            reply.SetStatus(NRTYServer::TReply::NOTNOW_ONSTOP);
        }
    } catch (...) {
        reply.SetStatus(NRTYServer::TReply::INTERNAL_ERROR);
        reply.SetStatusMessage(CurrentExceptionMessage());
    }
    AtomicDecrement(InFlight);

    return reply;
}

void NFusion::TBaseIndexingClient::Shutdown() {
    Active = false;
}

size_t NFusion::TBaseIndexingClient::GetQueueSize() const {
    return Queue.Size();
}

size_t NFusion::TBaseIndexingClient::GetInFlight() const {
    return InFlight;
}

void* NFusion::TBaseIndexingClient::CreateThreadSpecificResource() const {
    ThreadDisableBalloc();
    return nullptr;
}

void NFusion::TBaseIndexingClient::DestroyThreadSpecificResource(void*) const{
}

NFusion::TRemoteIndexingClient::TRemoteIndexingClient(const TString& host, ui16 port, const TIndexingClientOptions& options)
    : TBaseIndexingClient(NRTYServer::IndexingNehProtocol, host, port, options)
{
}

NFusion::TInternalIndexingClient::TInternalIndexingClient(ui16 port, const TIndexingClientOptions& options)
    : TBaseIndexingClient(NRTYServer::IndexingInternalProtocol, "localhost", port, options)
{
}
