#include "zoo_queue.h"
#include <library/cpp/zookeeper/zookeeper.h>
#include <library/cpp/regex/pcre/regexp.h>
#include <library/cpp/string_utils/base64/base64.h>
#include <util/folder/path.h>
#include <library/cpp/logger/global/global.h>
#include <util/generic/set.h>
#include <util/string/vector.h>
#include <util/string/cast.h>

namespace NRTYDeploy {

    TZooQueue::TZooQueue(const IVersionedStorage::TOptions& options) {
        IVersionedStorage::TOptions copy = options;
        copy.ZooOptions.Root = options.QueueName;
        copy.LocalOptions.Root = options.QueueName;
        copy.MongoOptions.CollectionPrefix = options.QueueName + "_";
        Storage.Reset(IVersionedStorage::Create(copy));
    }

    bool TZooQueue::Enqueue(const TString& queueName, const TString& subQueueName, const TString& id) const {
        try {
            DEBUG_LOG << "Enqueue queue=" << queueName << ", subQueue=" << subQueueName << ", id=" << id << "..." << Endl;
            TString pathLock = TFsPath(queueName + "/" + subQueueName).Fix().GetPath();
            TAbstractLock::TPtr lock = Storage->WriteLockNode(pathLock);
            if (!Storage->CreatePersistentSequentialNode("/queues/" + queueName + "/" + subQueueName + "/" + id, "")) {
                return false;
            }
            DEBUG_LOG << "Enqueue queue=" << queueName << ", subQueue=" << subQueueName << ", id=" << id << "...OK" << Endl;
        } catch (yexception& e) {
            DEBUG_LOG << "Can't enqueue: " << e.what() << " " << queueName << "-" << subQueueName << "-" << id << Endl;
            return false;
        }
        return true;
    }

    bool TZooQueue::Dequeue(const TString& queueName, TDequeueInfo& info) const {
        TString path = TFsPath("/queues/" + queueName).Fix().GetPath();
        TString subQueue;
        try {
            TVector<TString> nodes;
            TSet<TString> services;
            if (Storage->GetNodes(path, nodes, true)) {
                for (size_t i = 0; i < nodes.size(); ++i) {
                    subQueue = nodes[i];
                    TString pathLock = TFsPath(queueName + "/" + subQueue).Fix().GetPath();
                    info.Lock = Storage->WriteLockNode(pathLock, TDuration::Seconds(2));
                    if (!!info.Lock) {
                        DEBUG_LOG << subQueue << " locked" << Endl;
                        break;
                    } else {
                        DEBUG_LOG << subQueue << " locked already" << Endl;
                    }
                }
                if (!!info.Lock) {
                    nodes.clear();
                    if (!Storage->GetNodes(path + "/" + subQueue, nodes, true) || !nodes.size()) {
                        ERROR_LOG << "cannot get subnodes queue=" << queueName << ", subQueue=" << subQueue << Endl;
                        Storage->RemoveNode(path + "/" + subQueue, true);
                        info.Lock.Reset(nullptr);
                        return false;
                    }
                    DEBUG_LOG << "Dequeue queue=" << queueName << ", subQueue=" << subQueue << ", ids=" << nodes.size() << Endl;
                    ui32 minVersion = Max<ui32>();
                    for (ui32 i = 0; i < nodes.size(); ++i) {
                        ui32 ver = FromString(nodes[i].substr(nodes[i].size() - 10));
                        if (minVersion > ver) {
                            info.FullId = nodes[i];
                            info.Id = nodes[i].substr(0, nodes[i].size() - 10);
                            info.SubQueue = subQueue;
                            info.FullPath = path + "/" + subQueue + "/" + nodes[i];
                            minVersion = ver;
                        }
                    }
                    DEBUG_LOG << "Dequeue task id=" << info.FullId << Endl;
                }
            }
        } catch (yexception& e) {
            DEBUG_LOG << "Can't dequeue: " << e.what() << Endl;
            return false;
        }
        return !!info.Lock;
    }

    bool TZooQueue::GetTasks(const TString& queueName, const TString& subQueueRegEx, TVector<TString>& tasks) const {
        TRegExMatch rem(subQueueRegEx.data());
        tasks.clear();
        TVector<TString> subQueues;
        if (Storage->GetNodes("/queues/" + queueName, subQueues, true)) {
            for (ui32 i = 0; i < subQueues.size(); ++i) {
                if (rem.Match(subQueues[i].data())) {
                    TVector<TString> subTasks;
                    Storage->GetNodes("/queues/" + queueName + "/" + subQueues[i], subTasks, true);
                    tasks.insert(tasks.end(), subTasks.begin(), subTasks.end());
                }
            }
            return true;
        } else
            return false;
    }

    bool TZooQueue::Ack(const TDequeueInfo& info) const {
        DEBUG_LOG << "Remove task " << info.FullPath << "..." << Endl;
        TVector<TString> subnodes;
        TString parentPath = TFsPath(info.FullPath).Parent().GetPath();
        if (!Storage->GetNodes(parentPath, subnodes, true))
            ERROR_LOG << "cannot get subnodes " << parentPath << Endl;
        if (subnodes.size() == 1) {
            Storage->RemoveNode(parentPath);
            DEBUG_LOG << "Remove subQueue " << parentPath << "...OK" << Endl;
        } else {
            Storage->RemoveNode(info.FullPath);
            DEBUG_LOG << "Remove task " << info.FullPath << " (" << subnodes.size() << ")...OK" << Endl;
        }
        return true;
    }
}
