#include "lock.h"

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

#include <util/string/join.h>
#include <util/string/vector.h>


namespace NSaasLB {
    TConsumerLocker::TConsumerLocker(
        const TString& servers,
        const TString& path,
        const TString& topic,
        const TString& entityId,
        std::function<void()> connectionLostCallback,
        ui32 maxAttempts /*= 3*/,
        TDuration waitTime /*= TDuration::Minutes(5)*/
    )
        : TZKClient(servers, connectionLostCallback, maxAttempts, waitTime)
        , Path(path)
        , Topic(topic)
        , EntityId(entityId)
    {
        UseNewLBProtocol = Topic.Contains('/');
        Topic = EncodeTopic(Topic);
    }

    TConsumerLocker::~TConsumerLocker() {
        Release();
    }

    TString TConsumerLocker::DecodeConsumer(const TString& path) const {
        if (!UseNewLBProtocol) {
            return path;
        }
        TString decoded = path;
        std::replace(decoded.begin(), decoded.vend(), '@', '/');
        return decoded;
    }

    TString TConsumerLocker::EncodeTopic(const TString& path) const {
        if (!path.Contains("/")) {
            return path;
        }
        auto parts = SplitString(path, "/");
        TString name = parts.back();
        return Join("--", JoinRange("@", parts.begin(), parts.end() - 1), name);
    }

    TString TConsumerLocker::LockFreeConsumer() {
        if (LockedConsumer) {
            return LockedConsumer;
        }
        INFO_LOG << "ConsumerLocker: Searching for a free consumer..." << Endl;
        if (!IsConnected()) {
            Reconnect();
        }
        auto consumers = GetChildren(Path);
        for (const auto& consumer: consumers) {
            const TString& lockPath = Path + "/" + consumer + "/" + Topic;
            if (CreateLock(lockPath, EntityId)) {
                LockedConsumer = DecodeConsumer(consumer);
                INFO_LOG << "ConsumerLocker: Locked consumer=" << LockedConsumer << Endl;
                return LockedConsumer;
            }
        }
        INFO_LOG << "ConsumerLocker: No free consumers" << Endl;
        return "";
    }

    void TConsumerLocker::Release() {
        Disconnect();

        INFO_LOG << "ConsumerLocker: Released consumer lock" << Endl;
        LockedConsumer = "";
    }
};
