#pragma once

#include "settings.h"
#include "shards_info.h"

#include <saas/api/action.h>
#include <saas/library/check_message/check_message.h>
#include <saas/library/sharding/sharding.h>
#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h>

#include <library/cpp/threading/atomic/bool.h>

#include <util/thread/pool.h>

class TThread;

namespace NSaas {

class TPersQueueWriter {
private:
    class TMessageSerializer : protected TMessageChecker {
    public:
        TMessageSerializer(const TCheckMessageSettings& settings, NPersQueueCommon::ECodec codec);

        NPersQueue::TData CheckAndSerialize(const NRTYServer::TMessage& message);

    private:
        ui32 MaxSizeBytes;
        NPersQueueCommon::ECodec Codec;
    };

public:
    struct TShardProducer;
    using TShardProducerPtr = TIntrusivePtr<TShardProducer>;

    struct TWriteResult {
        struct TAttemtResult {
            TAttemtResult(ui16 attempt, bool written, TStringBuf shard, TStringBuf comment, ui32 size, ui32 compressedSize);
            NJson::TJsonValue ToJson() const;

            ui16 Attempt = 0;
            bool Written = false;
            TString Shard;
            TString Comment;
            ui32 Size = 0;
            ui32 CompressedSize = 0;
        };

        TWriteResult() = default;
        TWriteResult(bool written, bool userError, bool exceededInFlight, TStringBuf comment, const TVector<TAttemtResult>& attempts = {});
        NJson::TJsonValue ToJson() const;
        TString ToString() const;

        bool Written = false;
        bool UserError = false;
        bool ExceededInFlight = false;
        TString Comment;
        TVector<TAttemtResult> Attempts;
    };
    using TWriteResultFuture = NThreading::TFuture<TPersQueueWriter::TWriteResult>;
private:
    using TWriteFuture = NThreading::TFuture<NPersQueue::TProducerCommitResponse>;

    struct TWriteState;
    struct TWriteJob;
    struct TThreadNamer {
        void* CreateThreadSpecificResource() const;
        void DestroyThreadSpecificResource(void*) const;
        void SetName(const TString& name);
        const TString& GetName() const;
    private:
        TString Name;
    };
    using TWriteJobPtr = TIntrusivePtr<TWriteJob>;

private:
    TWriteResult::TAttemtResult GetAttemptInfo(const TWriteState& status) const;
    bool Write(TWriteJobPtr job, ui32 stateIndex);
    void Write(TWriteJobPtr job);
    void ProcessProducerResponse(TWriteJobPtr job, ui32 stateIndex, ui32 attempt);
    void ProcessDeadline(TWriteJobPtr job, ui32 stateIndex, ui32 attempt, TInstant deadline);

    TString CreateSourceId() const;
    ui64 GenerateMessageId() const;

    void PrepareMessage(NRTYServer::TMessage& message) const;
    bool MessageNeedsPrepare(const NRTYServer::TMessage& message) const;
    NPersQueue::TData GetPreparedData(const NRTYServer::TMessage& originalMessage) const;

    void InitThreads();
    THolder<IThreadPool> CreateThreadPool();

protected:
    virtual void BeforeProcessWrite(const NRTYServer::TMessage& msg, const NPersQueue::TData& data, const TVector<NSearchMapParser::TShardsInterval>& shards) const;
    virtual void InitShards(const TVector<NSearchMapParser::TShardsInterval>& shards) const;

public:
    TPersQueueWriter();
    virtual ~TPersQueueWriter();

    void Init(const TPersQueueWriterSettings& settings);
    void Stop();
    ui32 GetInFlight() const;
    NJson::TJsonValue GetStatus() const;
    NJson::TJsonValue GetJsonConfig() const;

    static TWriteResult GetWriteErrorInfo(TStringBuf comment, bool userError = false, bool exceededInFlight = false);

    TWriteResultFuture Write(const NRTYServer::TMessage& message);
    TWriteResultFuture Write(const NSaas::TAction& action);
    TWriteResultFuture Write(const NJson::TJsonValue& document);

    static TString GetTopicName(TStringBuf directoryWithTopics, TStringBuf shardName);
    static TString GetTopicName(
        TStringBuf directoryWithTopics,
        NSearchMapParser::TShardsInterval shard
    );

    void UpdateShards(const TServiceShards& serviceShards);
    void UpdateSearchMap(const NSearchMapParser::TSearchMap& searchMap);

private:
    TWriteResultFuture CreateWriteJob(const NRTYServer::TMessage& message);
    THashMap<TString, TShardProducerPtr> CreateProducers(const TVector<NSearchMapParser::TShardsInterval>& shards) const;
    void UpdateSearchMap() noexcept;
    static void* SearchMapUpdaterProc(void* object);

private:
    TPersQueueWriterSettings Settings;
    NPersQueue::TProducerSettings ProducerSettingsTemplate;

    THolder<TMessageSerializer> MessageSerializer;

    TRWMutex ShardsMutex;
    TServiceShards ServiceShards;
    THashMap<TString, TShardProducerPtr> ProducerPerShard;
    TInstant ShardsUpdateTime;

    TPQLibPtr PQLib;

    NAtomic::TBool Initilized = false;
    TAtomic InFlight = 0;

    TThreadNamer WriteThreadNamer;
    TThreadNamer WaitThreadNamer;
    TThreadNamer UpdateSearchmapThreadNamer;

    THolder<IThreadPool> WriteQueue;
    THolder<IThreadPool> WaitQueue;
    THolder<TThread> SearchMapUpdater;
    TManualEvent Stopped;

    THolder<ITelemetry> Telemetry;
};

}
