#include "zoo_storage.h"

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

#include <saas/util/logging/trace.h>

#include <util/folder/path.h>
#include <util/string/cast.h>
#include <util/string/subst.h>
#include <util/string/vector.h>
#include <util/system/guard.h>

namespace NSaas {

    IVersionedStorage::TFactory::TRegistrator<TZooStorage> TZooStorage::Registrator(IVersionedStorage::TOptions::TStorageType::ZOO);

    TZooLock::~TZooLock() {
        ReleaseLock();
    }

    bool TZooLock::BuildLock() {
        try {
            VERIFY_WITH_LOG(!LockName, "Incorrect usage");
            TString pathLocks = TFsPath("/locks/abstract_locks/" + Id).Fix().GetPath();
            Storage.MkDirs(pathLocks);
            TString lockPath = pathLocks + (WriteLock ? "/write-" : "/read-") + Id;
            LockName = Storage.GetZK()->Create(lockPath, "", NZooKeeper::ACLS_ALL, NZooKeeper::CM_EPHEMERAL_SEQUENTIAL);
        } catch (...) {
            ERROR_LOG << CurrentExceptionMessage() << Endl;
            return false;
        }
        return true;
    }

    bool TZooLock::IsLockTaken() const {
        VERIFY_WITH_LOG(!!LockName, "Incorrect usage");
        TString pathLocks = TFsPath("/locks/abstract_locks/" + Id).Fix().GetPath();
        ui32 version = FromString<ui32>(LockName.substr(LockName.size() - 10, 10));
        TVector<TString> locks = Storage.GetZK()->GetChildren(pathLocks, false);
        for (ui32 i = 0; i < locks.size(); ++i) {
            if (FromString<ui32>(locks[i].substr(locks[i].size() - 10, 10)) < version) {
                bool isWriteLock = locks[i].StartsWith("write-");
                if (WriteLock || isWriteLock)
                    return false;
            }
        }
        return true;
    }

    bool TZooLock::ReleaseLock() {
        if (!!LockName) {
            Storage.FastRemoveNode(LockName);
            LockName = "";
            return true;
        }
        return false;
    }

    bool TZooLock::IsLockedImpl() const {
        VERIFY_WITH_LOG(!!LockName, "Lock should be acquired via TryLock first");
        return Storage.ExistsNode(LockName);
    }

    bool TZooLock::IsLocked() const {
        return IsLocked_.Get();
    }

    TZooLock::TZooLock(const TZooStorage& storage, const TString& id, bool writeLock)
        : Storage(storage)
        , Id(id)
        , WriteLock(writeLock)
        , IsLocked_(*this)
    {
        SubstGlobal(Id, '/', '_');
        BuildLock();
    }

    TAbstractLock::TPtr TZooStorage::NativeLockNode(const TString& path, const bool isWrite, const TDuration timeout) const {
        TFsPath fs(path);
        TInstant timeStart = Now();
        TAutoPtr<TZooLock> result = new TZooLock(*this, fs.Fix().GetPath(), isWrite);
        while (timeStart + timeout > Now()) {
            try {
                if (result->IsLockTaken()) {
                    return result.Release();
                } else {
                    Sleep(TDuration::MilliSeconds(10));
                }
            } catch (NZooKeeper::TInvalidStateError&) {
                BuildZK("TZooLock::ReadLock failed: " + CurrentExceptionMessage());
            } catch (NZooKeeper::TConnectionLostError&) {
                BuildZK("TZooLock::ReadLock connection lost");
            }
        }
        return nullptr;
    }

    TAbstractLock::TPtr TZooStorage::NativeReadLockNode(const TString& path, TDuration timeout) const {
        return NativeLockNode(path, false, timeout);
    }

    TAbstractLock::TPtr TZooStorage::NativeWriteLockNode(const TString& path, TDuration timeout) const {
        return NativeLockNode(path, true, timeout);
    }

    bool TZooStorage::GetVersion(const TString& key, i64& version) const {
        try {
            TReadGuard wg(RWMutex);
            TString path = TFsPath("/" + key).Fix().GetPath();
            if (!!ZK->Exists(path)) {
                TString ver = ZK->GetData(path);
                if (!ver) {
                    return false;
                } else {
                    try {
                        version = FromString(ver);
                    } catch (...) {
                        return false;
                    }
                }
                return true;
            }
        } catch (NZooKeeper::TInvalidStateError&) {
            BuildZK("GetVersion failed: " + CurrentExceptionMessage());
            return GetVersion(key, version);
        } catch (NZooKeeper::TConnectionLostError&) {
            BuildZK("GetVersion failed: " + CurrentExceptionMessage());
            return GetVersion(key, version);
        }
        return false;

    }

    bool TZooStorage::BuildZK(const TString& reason) const {
        TWriteGuard wg(RWMutex);
        WARNING_LOG << "ZK rebuild on " << reason << " starting..." << Endl;
        for (ui32 i = 0; i < 100; ++i) {
            NZooKeeper::TZooKeeperOptions options(ZooAddresses + TFsPath("/" + Root).Fix().GetPath());
            options.Timeout = TDuration::MilliSeconds(10000000);
            options.LogLevel = static_cast<NZooKeeper::ELogLevel>(ZooLogLevel);
            ZK.Reset(new NZooKeeper::TZooKeeper(options));
            TVector<TString> nodes;
            try {
                ZK->Exists("/");
                INFO_LOG << "ZK rebuild finished" << Endl;
                return true;
            }
            catch (...) {
                WARNING_LOG << "Can't works with zookeeper : " << CurrentExceptionMessage() << Endl;
            }
            Sleep(TDuration::MilliSeconds(500));
        }
        FAIL_LOG("Can't use ZooKeeper %s for root = %s", ZooAddresses.data(), Root.data());
        return false;
    }

    TZooStorage::TZooStorage(const TOptions& options)
        : IVersionedStorage(options)
        , ZooAddresses(options.ZooOptions.Address)
        , Root("/" + options.ZooOptions.Root)
        , ZooLogLevel(options.ZooOptions.LogLevel)
    {
        BuildZK("starting");
    }

    bool TZooStorage::FastRemoveNode(const TString& key) const {
        try {
            TString path = TFsPath("/" + key).Fix().GetPath();
            TReadGuard wg(RWMutex);
            NZooKeeper::TStatPtr ex = ZK->Exists(path);
            if (!!ex)
                ZK->Delete(path, ex->version);
            return !!ex;
        }
        catch (NZooKeeper::TInvalidStateError&) {
            BuildZK("RemoveNode failed: " + CurrentExceptionMessage());
            return FastRemoveNode(key);
        }
        catch (NZooKeeper::TConnectionLostError&) {
            BuildZK("RemoveNode failed: " + CurrentExceptionMessage());
            return FastRemoveNode(key);
        }
        catch (NZooKeeper::TNodeNotExistsError&) {
            WARNING_LOG << "Race for node " << key << " removing" << Endl;
            return true;
        }
        return false;
    }

    bool TZooStorage::RemoveNode(const TString& key, bool withHistory) const {
        try {
            TString path = TFsPath("/" + key).Fix().GetPath();
            NZooKeeper::TStatPtr ex;
            {
                TReadGuard wg(RWMutex);
                ex = ZK->Exists(path);
            }
            if (!!ex) {

                TVector<TString> nodes;
                if (GetNodes(key, nodes, true)) {
                    for (ui32 i = 0; i < nodes.size(); ++i) {
                        RemoveNode("/" + key + "/" + nodes[i]);
                    }
                }

                for (ui32 part = 0;; ++part) {
                    TString pathPart = path + (part ? ".$" + ToString(part) : "");
                    if (FastRemoveNode(pathPart)) {
                        if (withHistory)
                            FastRemoveNode("/history/" + pathPart);
                    } else
                        return true;
                }
            } else {
                return false;
            }

        } catch (NZooKeeper::TInvalidStateError&) {
            BuildZK("RemoveNode failed: " + CurrentExceptionMessage());
            return RemoveNode(key, withHistory);
        } catch (NZooKeeper::TConnectionLostError&) {
            BuildZK("RemoveNode failed: " + CurrentExceptionMessage());
            return RemoveNode(key, withHistory);
        } catch (NZooKeeper::TNodeNotExistsError&) {
            WARNING_LOG << "Race for node " << key << " removing" << Endl;
            return true;
        }
        return false;
    }

    bool TZooStorage::ExistsNode(const TString& key) const {
        try {
            TReadGuard wg(RWMutex);
            TString path = TFsPath("/" + key).Fix().GetPath();
            return !!ZK->Exists(path);
        } catch (NZooKeeper::TInvalidStateError&) {
            BuildZK("ExistsNode failed: " + CurrentExceptionMessage());
            return ExistsNode(key);
        } catch (NZooKeeper::TConnectionLostError&) {
            BuildZK("ExistsNode failed: " + CurrentExceptionMessage());
            return ExistsNode(key);
        }
    }

    bool TZooStorage::GetValue(const TString& key, TString& result, i64 version, bool lock) const {
        NUtil::TTracer timer("get_value " + key);
        try {
            TString path;
            if (version == -1) {
                TAbstractLock::TPtr lockPtr = nullptr;
                if (lock) {
                    NUtil::TTracer lockTimer("lock " + key);
                    lockPtr = ReadLockNode(key);
                }
                if (GetVersion(key, version))
                    return GetValue(key, result, version);
                else {
                    WARNING_LOG << "Incorrect version " << version << " for " << key << Endl;
                    return false;
                }
            } else {
                TString ver;
                sprintf(ver, "%.10ld", version);
                path = TFsPath("/history/" + key + ver).Fix().GetPath();
            }
            TReadGuard rg(RWMutex);
            TValueCache::const_iterator cache = ValueCache.find(path);
            if (cache != ValueCache.end()) {
                result = cache->second;
                DEBUG_LOG << "Loaded from cache " << key << Endl;
                return true;
            }
            if (!!ZK->Exists(path)) {
                result = "";
                for (ui32 part = 0; ; ++part) {
                    TString pathPart = path + (part ? ".$" + ToString(part) : "");
                    NZooKeeper::TStatPtr stat = ZK->Exists(pathPart);
                    if (!stat) {
                        rg.Release();
                        TWriteGuard wg(RWMutex);
                        ValueCache[path] = result;
                        return true;
                    }
                    result += ZK->GetData(pathPart);
                }
            }
        } catch (NZooKeeper::TInvalidStateError&) {
            BuildZK("GetValue failed: " + CurrentExceptionMessage());
            return GetValue(key, result, version);
        } catch (NZooKeeper::TConnectionLostError&) {
            BuildZK("GetValue failed: " + CurrentExceptionMessage());
            return GetValue(key, result, version);
        }
        catch (NZooKeeper::TNodeNotExistsError&) {
            WARNING_LOG << "Race for node " << key << " get" << Endl;
            return false;
        }
        return false;
    }

    void TZooStorage::MkDirs(const TString& path) const {
        Y_VERIFY(path.StartsWith('/'), "path should be absolute");
        try {
            TFsPath fs(path);
            fs = fs.Fix();
            if (fs.GetPath() != "/")
                MkDirs(fs.Parent().GetPath());
            {
                TReadGuard wg(RWMutex);
                NZooKeeper::TStatPtr ex = ZK->Exists(fs.GetPath());
                if (!ex) {
                    ZK->Create(fs.GetPath(), "", NZooKeeper::ACLS_ALL, NZooKeeper::CM_PERSISTENT);
                }
            }
        } catch (NZooKeeper::TInvalidStateError&) {
            BuildZK("MkDirs failed: " + CurrentExceptionMessage());
            return MkDirs(path);
        } catch (NZooKeeper::TNodeExistsError&) {
            return;
        } catch (NZooKeeper::TConnectionLostError&) {
            BuildZK("MkDirs failed: " + CurrentExceptionMessage());
            return MkDirs(path);
        }
    }

    bool TZooStorage::SetValue(const TString& key, const TString& value, bool storeHistory, bool lock, i64* versionOut) const {
        TDebugLogTimer timer("set_value " + key);
        try
        {
            TString currentValue;
            TString pathHistory = TFsPath("/history/" + key).Fix().GetPath();
            MkDirs(TFsPath("/history/" + key).Fix().Parent().GetPath());
            TAbstractLock::TPtr lockPtr = nullptr;
            if (lock) {
                TDebugLogTimer lockTimer("lock " + key);
                lockPtr = WriteLockNode(key);
            }
            TString startPath;
            for (ui32 part = 0;; ++part) {
                if (720 * 1024 * part >= value.size() && part)
                    break;
                TReadGuard wg(RWMutex);
                if (!part)
                    startPath = ZK->Create(pathHistory, value.substr(part * 720 * 1024, 720 * 1024), NZooKeeper::ACLS_ALL, NZooKeeper::CM_PERSISTENT_SEQUENTIAL);
                else
                    ZK->Create(startPath + ".$" + ToString(part), value.substr(part * 720 * 1024, 720 * 1024), NZooKeeper::ACLS_ALL, NZooKeeper::CM_PERSISTENT);
            }
            TString ver;
            ver = startPath.substr(startPath.size() - 10, TString::npos);
            i32 version = FromString<i32>(ver);
            if (versionOut)
                *versionOut = version;
            TString path = TFsPath("/" + key).Fix().GetPath();
            {
                i64 oldVersion = -1;
                if (!storeHistory)
                    GetVersion(path, oldVersion);
                if (version < oldVersion)
                    return false;
                TReadGuard wg(RWMutex);
                if (!!ZK->Exists(path)) {
                    ZK->SetData(path, ver);
                    if (!storeHistory && oldVersion > 0) {
                        TString verOld;
                        sprintf(verOld, "%.10ld", oldVersion);
                        TString pathDelHistory = TFsPath("/history/" + key + verOld).Fix().GetPath();
                        for (ui32 part = 0;; ++part) {
                            TString pathPart = pathDelHistory + (part ? ".$" + ToString(part) : "");
                            NZooKeeper::TStatPtr stat = ZK->Exists(pathPart);
                            if (!stat)
                                break;
                            ZK->Delete(pathPart, stat->version);
                        }
                    }
                    return true;
                }
            }
            MkDirs(TFsPath("/" + key).Fix().Parent().GetPath());
            {
                TReadGuard wg(RWMutex);
                ZK->Create(path, ver, NZooKeeper::ACLS_ALL, NZooKeeper::CM_PERSISTENT);
            }
            return true;
        } catch (NZooKeeper::TInvalidStateError&) {
            BuildZK("SetValue failed: " + CurrentExceptionMessage());
            return SetValue(key, value);
        } catch (NZooKeeper::TConnectionLostError&) {
            BuildZK("SetValue failed: " + CurrentExceptionMessage());
            return SetValue(key, value);
        }
        return false;
    }

    bool TZooStorage::GetNodes(const TString& key, TVector<TString>& result, bool withDirs) const {
        TString path = TFsPath("/" + key).Fix().GetPath();
        try {
            TReadGuard wg(RWMutex);
            if (!ZK->Exists(path))
                return false;
            result = ZK->GetChildren(path);
            if (!withDirs) {
                for (int i = result.size() - 1; i >= 0; i--) {
                    try{
                        if (!ZK->GetData(path + "/" + result[i])) {
                            result.erase(result.begin() + i);
                        }
                    } catch (...){
                        result.erase(result.begin() + i);
                        DEBUG_LOG << "Unknown error, path " << path << Endl;
                    }
                }
            }
            for (int i = result.size() - 1; i >= 0; i--) {
                if (result[i].find(".$") != TString::npos) {
                    result.erase(result.begin() + i);
                }
            }
            return true;
        } catch (NZooKeeper::TInvalidStateError&) {
            BuildZK("GetNodes failed: " + CurrentExceptionMessage());
            return GetNodes(key, result, withDirs);
        } catch (NZooKeeper::TConnectionLostError&) {
            BuildZK("GetNodes failed: " + CurrentExceptionMessage());
            return GetNodes(key, result, withDirs);
        } catch (NZooKeeper::TNodeNotExistsError&) {
            WARNING_LOG << "Race for node " << key << " get child" << Endl;
            return true;
        }
        return false;
    }

    bool TZooStorage::CreatePersistentSequentialNode(const TString& key, const TString& data) const {
        try {
            TFsPath path = TFsPath("/" + key).Fix();
            MkDirs(path.Parent());
            GetZK()->Create(path, data, NZooKeeper::ACLS_ALL, NZooKeeper::CM_PERSISTENT_SEQUENTIAL);
            return true;
        } catch (...) {
            ERROR_LOG << "Exception on node " << key << " construction: " << CurrentExceptionMessage() << Endl;
            return false;
        }
    }

}
