#pragma once

#include <optional>

#include <library/cpp/logger/global/global.h>
#include <library/cpp/zookeeper/zookeeper.h>

#include <util/system/rwlock.h>
#include <util/system/thread.h>


namespace NSaas {
    class TZKClient : public NZooKeeper::IWatcher {
    private:
        NZooKeeper::TZooKeeperOptions Options;
        std::function<void()> ConnectionLostCallback;
        ui32 MaxAttempts;
        TDuration WaitTime;

        THolder<NZooKeeper::TZooKeeper> Zookeeper;
        TRWMutex RWMutex;

        TManualEvent ConnectionEvent;
        TAtomic LastEventState = NZooKeeper::KS_DISCONNECTED;
        THolder<TThread> WaitConnectionThread;

    protected:
        bool CreateLock(const TString& path, const TString& entityId);
        TVector<TString> GetChildren(const TString& path);
        TString GetData(const TString& path);

        void Reconnect();
        bool IsConnected();
        void Disconnect();

        template <class T>
        std::optional<T> ExecuteWithRetries(const std::function<T()>& func, const TString& errorMessage) {
            for (ui32 attempt = 0; attempt < MaxAttempts; ++attempt) {
                try {
                    TReadGuard wg(RWMutex);
                    if (!Zookeeper) {
                        ythrow yexception() << "TZKClient: Zookeeper uninitialized";
                    }
                    return func();
                } catch (NZooKeeper::TInvalidStateError&) {
                    Reconnect();
                } catch (NZooKeeper::TConnectionLostError&) {
                    ConnectionEvent.WaitT(WaitTime);
                } catch (NZooKeeper::TNodeExistsError&) {
                    return {};
                } catch (...) {
                    ERROR_LOG << "TZKClient: " << errorMessage << ", error: " << CurrentExceptionMessage() << Endl;
                    Reconnect();
                }
            }
            return {};
        }

    public:
        TZKClient(
            const TString& servers,
            std::function<void()> connectionLostCallback = nullptr,
            ui32 MaxAttempts = 3,
            TDuration waitTime = TDuration::Minutes(5)
        );
        virtual ~TZKClient();

        void Process(const NZooKeeper::TWatchedEvent& event) override;
    };
}
