#pragma once
#include "abstract_context.h"
#include "sender_context.h"

#include <saas/indexerproxy/unistat_signals/signals.h>
#include <saas/indexerproxy/configs/config.h>
#include <saas/library/persqueue/logger/logger.h>
#include <saas/library/indexer_protocol/sender_neh.h>
#include <saas/library/behaviour/behaviour.h>
#include <saas/protos/rtyserver.pb.h>
#include <saas/library/persqueue/writer/shards_info.h>

#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h>

class TProxyConfig;
class TDispatchServiceSelector;

class THandlerOnSend;

class TPersQueueManager {
private:
    using TPQLibPtr = TAtomicSharedPtr<NPersQueue::TPQLib>;

private:
    class TShardsInfo : public NSaas::TServiceShards {
    public:
        using NSaas::TServiceShards::Init;
        void Init(ui32 shardsCount, NSaas::TShardsDispatcher::TPtr shardsDispatcher);
        TVector<TString> GetShardNames() const;
        TVector<TString> GetShardNames(const NRTYServer::TMessage& message) const;
    private:
        bool NumeratedShardsNames = false;
    };

    class TServiceWriter {
    private:
        class TWriter {
        private:
            THolder<NPersQueue::IProducer> Producer;
            NPersQueue::TProducerSeqNo LastSeqNo;
            TSpinLock Lock;
        public:
            TWriter(THolder<NPersQueue::IProducer>& producer, NPersQueue::TProducerSeqNo seqNo)
                : Producer(std::move(producer))
                , LastSeqNo(seqNo)
            {
            }
            NThreading::TFuture<NPersQueue::TProducerCommitResponse> Write(const NPersQueue::TData& data) {
                TGuard<TSpinLock> guard(Lock);
                NPersQueue::TProducerSeqNo seqNo = ++LastSeqNo;
                return Producer->Write(seqNo, data);
            }
        };
        using TWriterPtr = TAtomicSharedPtr<TWriter>;
        using TWriters = THashMap<TString, TWriterPtr>;

    private:
        TString GetTopicName(TStringBuf shard) const;
        TString CreateSourceId() const;
        void CreateWriter(TStringBuf shard);
        TWriterPtr GetWriter(TStringBuf shard);
        void IncrementInFlight(TStringBuf shard);
        void DecrementInFlight(TStringBuf shard);
        TVector<TString> GetShards(const NRTYServer::TMessage& message) const;
        NJson::TJsonValue DoSend(TStringBuf shard, const NPersQueue::TData& data);

    public:
        TServiceWriter(
            TPQLibPtr pqLib,
            std::shared_ptr<NPersQueue::ICredentialsProvider> credentialsProvider,
            const TString& service,
            const NSearchMapParser::TServiceSpecificOptions* serviceConfig,
            const TProxyServiceConfig& writeConfig,
            const NSearchMapParser::TSearchMap& searchMap,
            TIntrusivePtr<NPersQueue::ILogger> logger
        );
        bool DoSend(const NPersQueue::TData& data, TStringBuf shard, NJson::TJsonValue& attempts);
        void DoSend(const NRTYServer::TMessage& message, const THandlerOnSend* handler);

    private:
        TPQLibPtr PQLib;
        std::shared_ptr<NPersQueue::ICredentialsProvider> CredentialsProvider;
        const TString Service;
        TWriters WritersPerShard;
        const TString Server;
        const TString Ident;
        const TString SessionPrefix;
        const TDuration ConnectionTimeout;
        const TDuration SendTimeout;
        const ui32 SendAttemptsCount;
        TShardsInfo ShardsInfo;
        THashMap<TString, TAtomicCounter> InFlightPerShard;
        TIntrusivePtr<NPersQueue::ILogger> Logger;
        THolder<TPersqueueServiceSignals> Signals;
        TMutex Lock;
    };
    using TServiceWriterPtr = TAtomicSharedPtr<TServiceWriter>;
private:
    TPQLibPtr PQLib;
    TIntrusivePtr<NPersQueue::ILogger> Logger;
    THashMap<TString, TServiceWriterPtr> Services;
private:
    void CreateServiceWriter(const TString& service, std::shared_ptr<NPersQueue::ICredentialsProvider> credentialsProvider, const TProxyConfig& globalConfig);
    void CreateWriters(const TList<TString> services, const TProxyConfig& config);
public:
    void Init(const TProxyConfig& config);
    void DoSend(const TString& serviceName, const NRTYServer::TMessage& message, const THandlerOnSend* handler);
};

class TPersQueueReplier: public TReplier {
protected:
    virtual NRTYServer::TReply::TDispStatus DoGetDispStatus() const override;
    virtual bool DoSend(const TProxyConfig& config, TDispatchServiceSelector& selector) override;
    virtual bool DoVerify(const TProxyConfig& config, const TRTYMessageBehaviour& bh) const override;
    virtual void DoFinish() override;
    virtual bool DoStoreStatus(const TString& backend, const TString& status) override;

public:
    TPersQueueReplier(const TString& serviceName, ISenderTask& task, const NRTYServer::TMessage& message, const TDispatcherConfig* config, TDeferredMessagesQueue& storage);
};
