#include "mongo_storage.h"

#include <saas/util/mongo/connection.h>
#include <saas/util/mongo/exception.h>
#include <saas/util/mongo/engine.h>
#include <saas/util/mongo/pool.h>
#include <saas/util/logging/trace.h>

#include <library/cpp/json/json_value.h>

#include <util/generic/hash.h>
#include <util/generic/set.h>
#include <util/string/vector.h>
#include <library/cpp/string_utils/base64/base64.h>
#include <util/system/hostname.h>
#include <util/system/getpid.h>
#include <util/string/split.h>

namespace NSaas {

    IVersionedStorage::TFactory::TRegistrator<TMongoStorage> TMongoStorage::Registrator(IVersionedStorage::TOptions::TStorageType::MONGO);

    namespace {
        TString NormalizeKey(const TString& key) {
            return TFsPath("/" + key).Fix().GetPath();
        }
    }

    TMongoStorage::TMongoStorage(const IVersionedStorage::TOptions& options)
        : IVersionedStorage(options)
        , Config(options.MongoOptions)
        , ConnectionsPool(NUtil::TMongoEngine::CreateConnectionsPool(Config.Uri))
    {
        if (!Config.DB)
            ythrow yexception() << "there is no mongo db";
        IVersionedStorage::TOptions zooOpt = options;
        zooOpt.Type = IVersionedStorage::TOptions::ZOO;
        ZooStorage.Reset(IVersionedStorage::Create(zooOpt));
    }

    TMongoStorage::~TMongoStorage() {
    }

    TAbstractLock::TPtr TMongoStorage::NativeWriteLockNode(const TString& path, TDuration timeout) const {
        return ZooStorage->WriteLockNode(path, timeout);
    }

    TAbstractLock::TPtr TMongoStorage::NativeReadLockNode(const TString& path, TDuration timeout) const {
        return ZooStorage->ReadLockNode(path, timeout);
    }

    bool TMongoStorage::IsListNode(const TString& key) const {
        NJson::TJsonValue query(NJson::JSON_MAP);
        query["key"] = key;
        auto connection = ConnectionsPool->GetConection();
        return connection->GetCollection(Config.DB, Config.GetActualCollection()).Find(query, "key")->IsValid();
    }

    bool TMongoStorage::GetVersion(const TString& key, i64& version) const {
        TString normalizedKey = NormalizeKey(key);
        auto connection = ConnectionsPool->GetConection();
        NJson::TJsonValue query(NJson::JSON_MAP);
        query["key"] = normalizedKey;

        try {
            auto it = connection->GetCollection(Config.DB, Config.GetActualCollection()).Find(query, "version");
            if (!it->IsValid()) {
                return false;
            }
            version = it->Get()["version"].GetInteger();
        } catch (const NUtil::TMongoException& e) {
            ERROR_LOG << e.what() << "for key " << normalizedKey << Endl;
            return false;
        }

        return true;
    }

    bool TMongoStorage::RemoveNode(const TString& key, bool withHistory) const {
        TString normalizedKey = NormalizeKey(key);
        NJson::TJsonValue query(NJson::JSON_MAP);
        if (IsListNode(normalizedKey)) {
            query["key"] = normalizedKey;
        } else {
            query["key"] = NJson::TJsonValue(NJson::JSON_MAP);
            query["key"]["$regex"] = "^" + normalizedKey + (normalizedKey == "/" ? "" : "/");
        }

        auto connection = ConnectionsPool->GetConection();
        try {
            connection->GetCollection(Config.DB, Config.GetActualCollection()).Remove(query, true);
        } catch (const NUtil::TMongoException& e) {
            ERROR_LOG << "Can't remove data for key " << normalizedKey << ": " << e.what() << Endl;
            return false;
        }

        if (withHistory) {
            try {
                connection->GetCollection(Config.DB, Config.GetHistoryCollection()).Remove(query, true);
            } catch (const NUtil::TMongoException& e) {
                ERROR_LOG << "Can't remove history for key " << normalizedKey << ": " << e.what() << Endl;
                return false;
            }
        }
        return true;
    }

    bool TMongoStorage::ExistsNode(const TString& key) const {
        TString normalizedKey = NormalizeKey(key);
        if (IsListNode(normalizedKey)) {
            return true;
        }

        NJson::TJsonValue query(NJson::JSON_MAP);
        query["key"] = NJson::TJsonValue(NJson::JSON_MAP);
        query["key"]["$regex"] = "^" + normalizedKey + (normalizedKey == "/" ? "" : "/");

        auto connection = ConnectionsPool->GetConection();
        try {
            auto it = connection->GetCollection(Config.DB, Config.GetActualCollection()).Find(query, "key");
            return it->IsValid();
        } catch (const NUtil::TMongoException& e) {
            ERROR_LOG << e.what() << "for key " << normalizedKey << Endl;
        }
        return false;
    }

    bool TMongoStorage::GetValue(const TString& key, TString& result, i64 version, bool lock) const {
        TString normalizedKey = NormalizeKey(key);
        TAbstractLock::TPtr lockPtr = nullptr;
        if (lock) {
            TDebugLogTimer lockTimer("lock " + normalizedKey);
            lockPtr = ZooStorage->ReadLockNode("lock" + normalizedKey);
        }

        NJson::TJsonValue query(NJson::JSON_MAP);
        query["key"] = normalizedKey;
        if (version != LATEST_VERSION)
            query["version"] = version;

        auto connection = ConnectionsPool->GetConection();
        try {
            NUtil::TMongoCollection::TIterator::TPtr it;
            if (version == LATEST_VERSION) {
                it = connection->GetCollection(Config.DB, Config.GetActualCollection()).Find(query);
            } else {
                it = connection->GetCollection(Config.DB, Config.GetHistoryCollection()).Find(query);
            }

            if (!it->IsValid())
                return false;

            result = Base64Decode(it->Get()["value"].GetString());
        } catch (const NUtil::TMongoException& e) {
            ERROR_LOG << e.what() << "for key " << normalizedKey  << "(v:" << version << ")"<< Endl;
            return false;
        }

        return true;
    }

    namespace {
        TAtomic versionsCounter = 0;
    }

    bool TMongoStorage::SetValue(const TString& key, const TString& value, bool storeHistory, bool lock, i64* versionOut) const {
        TString normalizedKey = NormalizeKey(key);
        TAbstractLock::TPtr lockPtr = nullptr;
        if (lock) {
            TDebugLogTimer lockTimer("lock " + normalizedKey);
            lockPtr = ZooStorage->WriteLockNode("lock" + normalizedKey);
        }
        TString threadHash = HostName() + ":" + ToString(GetPID()) + ":" + ToString(AtomicIncrement(versionsCounter));
        i64 version = TInstant::Now().MilliSeconds();
        VERIFY_WITH_LOG(version < (Max<i64>() >> 16), "Version too big, it may become cause of problems");
        version = (version << 16) | (ComputeHash(threadHash) & (1 << 16 - 1));
        if (versionOut)
            *versionOut = version;
        NJson::TJsonValue document(NJson::JSON_MAP);
        document["key"] = normalizedKey;
        document["value"] = Base64Encode(value);
        document["version"] = version;

        NJson::TJsonValue query(NJson::JSON_MAP);
        query["key"] = normalizedKey;

        auto connection = ConnectionsPool->GetConection();
        try {
            connection->GetCollection(Config.DB, Config.GetHistoryCollection()).Insert(document);
            connection->GetCollection(Config.DB, Config.GetActualCollection()).Update(query, document);
        } catch (const NUtil::TMongoException& e) {
            ERROR_LOG << "Can't set value for key " << normalizedKey << ": " << e.what() << Endl;
            return false;
        }

        if (!storeHistory) {
            NJson::TJsonValue query(NJson::JSON_MAP);
            query["key"] = normalizedKey;
            query["version"] = NJson::TJsonValue(NJson::JSON_MAP);
            query["version"]["$ne"] = version;
            try {
                connection->GetCollection(Config.DB, Config.GetHistoryCollection()).Remove(query, true);
            } catch (const NUtil::TMongoException& e) {
                ERROR_LOG << "Can't remove history for key " << normalizedKey << ": " << e.what() << Endl;
                return false;
            }
        }

        return true;
    }

    bool TMongoStorage::GetNodes(const TString& key, TVector<TString>& result, bool withDirs) const {
        result.clear();
        TString normalizedKey = NormalizeKey(key);

        NJson::TJsonValue query(NJson::JSON_MAP);
        query["key"] = NJson::TJsonValue(NJson::JSON_MAP);
        query["key"]["$regex"] = "^" + normalizedKey + (normalizedKey == "/" ? "" : "/");

        auto connection = ConnectionsPool->GetConection();
        try {
            TSet<TString> dirs;
            auto it = connection->GetCollection(Config.DB, Config.GetActualCollection()).Find(query, "key");

            if (!it->IsValid()) {
                return false;
            }

            while (it->IsValid()) {
                TString oneKey(it->Get()["key"].GetStringRobust());

                TVector<TString> parts;
                StringSplitter(oneKey.data() + normalizedKey.size()).Split('/').SkipEmpty().Collect(&parts);

                if (parts.empty()) {
                    it->Next();
                    continue;
                }

                if (parts.size() > 1) {
                    if (withDirs)
                        dirs.insert(parts[0]);
                } else {
                    result.push_back(parts[0]);
                }

                it->Next();
            }
            if (withDirs)
                result.insert(result.end(), dirs.begin(), dirs.end());
        } catch (const NUtil::TMongoException& e) {
            ERROR_LOG << e.what() << "for key " << normalizedKey << Endl;
            return false;
        }

        return true;
    }

    bool TMongoStorage::CreatePersistentSequentialNode(const TString& key, const TString& data) const {
        TString normalizedKey = NormalizeKey(key);
        return SetValue(normalizedKey + Sprintf("%010lu", Now().MilliSeconds() % 1000000000), data, false);
    }

}
