#include "zk.h"

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

#include <util/string/vector.h>


namespace NSaas {
    TZKClient::TZKClient(
        const TString& servers,
        std::function<void()> connectionLostCallback,
        ui32 maxAttempts /*= 3*/,
        TDuration waitTime /*= TDuration::Minutes(5)*/
    )
        : Options(servers)
        , ConnectionLostCallback(connectionLostCallback)
        , MaxAttempts(maxAttempts)
        , WaitTime(waitTime)
    {
        Options.Watcher = this;
    }

    TZKClient::~TZKClient() {
        Disconnect();
    }

    bool TZKClient::CreateLock(const TString& path, const TString& entityId) {
        const auto& func = [&] () {
            Zookeeper->Create(path, entityId, NZooKeeper::ACLS_ALL, NZooKeeper::CM_EPHEMERAL);
            return true;
        };
        const TString& errorMessage = "Unable to create a lock on path=" + path;
        const auto& result = ExecuteWithRetries<bool>(func, errorMessage);
        return result.value_or(false);
    }

    TVector<TString> TZKClient::GetChildren(const TString& path) {
        const auto& func = [&]() {
            return Zookeeper->GetChildren(path);
        };
        const TString& errorMessage = "Unable to get children for path=" + path;
        const auto& result = ExecuteWithRetries<TVector<TString>>(func, errorMessage);
        if (result.has_value()) {
            return result.value();
        }
        return {};
    }

    TString TZKClient::GetData(const TString& path) {
        const auto& func = [&]() {
            return Zookeeper->GetData(path);
        };
        const TString& errorMessage = "Unable to read data by path=" + path;
        const auto& result = ExecuteWithRetries<TString>(func, errorMessage);
        return result.value_or("");
    }

    void TZKClient::Reconnect() {
        INFO_LOG << "TZKClient: Trying to reconnect..." << Endl;
        for (ui32 attempt = 0; attempt < MaxAttempts; ++attempt) {
            try {
                TWriteGuard wg(RWMutex);
                Zookeeper.Reset(new NZooKeeper::TZooKeeper(Options));
                return;
            } catch (NZooKeeper::TInvalidStateError&) {
            } catch (NZooKeeper::TConnectionLostError&) {
            } catch (...) {
                ERROR_LOG << "TZKClient: Unable to connect, error: " << CurrentExceptionMessage() << Endl;
            }
            {
                TWriteGuard wg(RWMutex);
                Zookeeper.Reset(nullptr);
            }
            ConnectionEvent.WaitT(WaitTime);
        }
        ERROR_LOG << "TZKClient: Unable to reconnect" << Endl;
    }

    bool TZKClient::IsConnected() {
        return (bool) Zookeeper;
    }

    void TZKClient::Disconnect() {
        if (!IsConnected()) {
            return;
        }
        TWriteGuard wg(RWMutex);
        Zookeeper.Reset();
        ConnectionEvent.Signal();
        WaitConnectionThread.Reset();
    }

    void TZKClient::Process(const NZooKeeper::TWatchedEvent& event) {
        if (event.KeeperState == AtomicGet(LastEventState)) {
            return;
        }
        AtomicSet(LastEventState, event.KeeperState);

        if (event.KeeperState == NZooKeeper::KS_EXPIRED) {
            ERROR_LOG << "TZKClient: The session has expired" << Endl;
            if (ConnectionLostCallback) {
                ConnectionLostCallback();
            }
        }

        if (event.KeeperState == NZooKeeper::KS_NO_SYNC_CONNECTED && !WaitConnectionThread) {
            INFO_LOG << "TZKClient: The connection has been lost, waiting the restoring..." << Endl;
            ConnectionEvent.Reset();
            TThread::TParams waitConnectionThreadParams([](void* object) -> void* {
                ThreadDisableBalloc();
                auto this_ = reinterpret_cast<TZKClient*>(object);
                if (!this_->ConnectionEvent.WaitT(this_->WaitTime)) {
                    ERROR_LOG << "TZKClient: The connection was lost after the wait time" << Endl;
                    if (this_->ConnectionLostCallback) {
                        this_->ConnectionLostCallback();
                    }
                }
                return nullptr;
            }, this);
            WaitConnectionThread.Reset(new TThread(waitConnectionThreadParams));
            WaitConnectionThread->Start();
        }

        if (event.KeeperState == NZooKeeper::KS_SYNC_CONNECTED) {
            INFO_LOG << "TZKClient: The connection has been establish" << Endl;
            ConnectionEvent.Signal();
            WaitConnectionThread.Destroy();
        }
    }
}
