#include "locks.h"

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

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

#include <algorithm>

namespace NSaasLB {

    TZooKeeperLocks::TZooKeeperLocks(const TString& servers)
        : Servers(servers)
    {
        Storage = MakeHolder<NZooKeeper::TZooKeeper>(NZooKeeper::TZooKeeperOptions(servers));
    }

    TZooKeeperLocks::~TZooKeeperLocks() = default;

    TVector<TString> TZooKeeperLocks::GetLocks(const TString& path) {
        auto& locks = LoadLocksByPath(path);
        return {locks.begin(), locks.end()};
    }

    void TZooKeeperLocks::Add(const TString& path, const TString& resource) {
        auto& locks = LoadLocksByPath(path);
        bool notExist = locks.insert(resource).second;
        if (notExist) {
            AddRequest(EChangeAction::CREATE, path, resource);
        }
    }

    void TZooKeeperLocks::Remove(const TString& path, const TString& resource) {
        auto& locks = LoadLocksByPath(path);
        auto it = locks.find(resource);
        if (it != locks.end()) {
            locks.erase(it);
            AddRequest(EChangeAction::REMOVE, path, resource);
        }
    }

    void TZooKeeperLocks::AddRequest(EChangeAction action, const TString& path, const TString& resource) {
        auto request = Changes.AddRequest();
        request->SetAction(action);
        auto lock = request->MutableLock();
        lock->SetResource(resource);
        lock->MutableZooKeeperLock()->SetNodePath(GetNodePath(path, resource));
    }

    THashSet<TString>& TZooKeeperLocks::LoadLocksByPath(const TString& path) {
        auto it = LocksByPath.find(path);
        if (it != LocksByPath.end()) {
            return it->second;
        }

        TVector<TString> resources = LoadItemsByPath(path);
        auto& locks = LocksByPath[path];
        for (auto& resource : resources) {
            locks.insert(resource);
        }
        return locks;
    }

    TVector<TString> TZooKeeperLocks::LoadItemsByPath(const TString& path) {
        TVector<TString> items;
        try {
            items = Storage->GetChildren(path);
        } catch (NZooKeeper::TNodeNotExistsError&) {
            return {};
        }

        for (auto& item : items) {
            item = Decode(item);
        }
        return items;
    }

    void TZooKeeperLocks::RecursiveCreateNode(const TString& nodePath) {
        auto dirs = SplitString(nodePath, "/");
        TString prefix = "";
        for (auto&& dir : dirs) {
            prefix += "/" + dir;
            if (!Storage->Exists(prefix)) {
                Storage->Create(prefix, "", NZooKeeper::ACLS_ALL, NZooKeeper::CM_PERSISTENT);
            }
        }
    }

    void TZooKeeperLocks::RecursiveRemoveNode(const TString& nodePath) {
        auto children = Storage->GetChildren(nodePath);
        for (auto& child : children) {
            RecursiveRemoveNode(Join("/", nodePath, child));
        }
        Storage->Delete(nodePath, -1);
    }

    TString TZooKeeperLocks::GetNodePath(const TString& path, const TString& resource) {
        return Join("/", path, Encode(resource));
    }

    TString TZooKeeperLocks::Encode(const TString& resource) const {
        TString encodedresource = resource;
        std::replace(encodedresource.begin(), encodedresource.vend(), '/', '@');
        return encodedresource;
    }

    TString TZooKeeperLocks::Decode(const TString& resource) const {
        TString decodedresource = resource;
        std::replace(decodedresource.begin(), decodedresource.vend(), '@', '/');
        return decodedresource;
    }

    TLockDescription TZooKeeperLocks::GetDescribed(const TString& path, const TString& resource) {
        auto& locks = LoadLocksByPath(path);
        Y_VERIFY(locks.count(resource));

        TString node = GetNodePath(path, resource);
        TLockDescription description;
        auto lock = description.MutableLock();
        lock->SetResource(resource);
        lock->MutableZooKeeperLock()->SetNodePath(node);

        if (Storage->Exists(node)) {
            auto topics = LoadItemsByPath(node);
            Sort(topics.begin(), topics.end(),[](const TString& a, const TString& b) {
                return (a.size() == b.size()) ? (a < b) : (a.size() < b.size());
            });
            for (auto& topic : topics) {
                TString topicLockNode = GetNodePath(node, topic);
                auto t = description.AddLockedTopics();
                t->SetTopic(topic);
                t->SetLockedBy(Storage->GetData(topicLockNode));
            }
        }
        return description;
    }

    const TLocksModifyRequests& TZooKeeperLocks::GetChanges() const {
        return Changes;
    }

    void TZooKeeperLocks::Apply(ui32 changesCount) {
        Y_ENSURE(changesCount <= ui32(Changes.GetRequest().size()), yexception() << "Incorrect changes count to apply");
        if (changesCount == 0) {
            changesCount = Changes.GetRequest().size();
        }
        INFO_LOG << "Apply locks changes " << changesCount << "/" << Changes.GetRequest().size() << Endl;
        for (ui32 i = 0; i < changesCount; ++i) {
            auto& request = Changes.GetRequest(i);
            auto nodePath = request.GetLock().GetZooKeeperLock().GetNodePath();
            switch (request.GetAction()) {
                case EChangeAction::CREATE: {
                    RecursiveCreateNode(nodePath);
                    break;
                }
                case EChangeAction::REMOVE: {
                    RecursiveRemoveNode(nodePath);
                    break;
                }
                default: {
                    ythrow yexception() << "unsupported action for zookeeper lock: " << request;
                }
            }
        }
        Changes.MutableRequest()->ExtractSubrange(0, changesCount, nullptr);
    }
}
