#include <saas/library/storage/distributed_lock.h>
#include <saas/library/storage/local/local_storage.h>

#include <library/cpp/testing/unittest/registar.h>
#include <library/cpp/logger/global/global.h>
#include <util/thread/factory.h>
#include <util/thread/pool.h>

#include <atomic>

void Clear(const TFsPath& path) {
    path.ForceDelete();
    path.MkDirs();
}

class TDisconnectableStorage : public NSaas::IVersionedStorage {
public:
    TDisconnectableStorage(NSaas::IVersionedStorage& slave)
        : IVersionedStorage(slave.GetOptions())
        , Slave(slave)
        , Disconnected(false)
    {}

    void Disconnect() {
        Disconnected = true;
    }

    bool RemoveNode(const TString& key, bool withHistory = false) const override {
        if (Disconnected)
            return false;
        return Slave.RemoveNode(key, withHistory);
    }

    bool ExistsNode(const TString& key) const override {
        if (Disconnected)
            return false;
        return Slave.ExistsNode(key);
    }

    bool GetVersion(const TString& key, i64& version) const override {
        if (Disconnected)
            return false;
        return Slave.GetVersion(key, version);
    }

    bool GetNodes(const TString& key, TVector<TString>& result, bool withDirs = false) const override {
        if (Disconnected)
            return false;
        return Slave.GetNodes(key, result, withDirs);
    }

    bool GetValue(const TString& key, TString& result, i64 version = -1, bool lock = true) const override {
        if (Disconnected)
            return false;
        return Slave.GetValue(key, result, version, lock);
    }

    bool SetValue(const TString& key, const TString& value, bool storeHistory = true, bool lock = true, i64* version = nullptr) const override {
        if (Disconnected)
            return false;
        return Slave.SetValue(key, value, storeHistory, lock, version);
    }

    bool CreatePersistentSequentialNode(const TString& key, const TString& data) const override {
        if (Disconnected)
            return false;
        return Slave.CreatePersistentSequentialNode(key, data);
    }

    NSaas::TAbstractLock::TPtr WriteLockNode(const TString& /*path*/, TDuration /*timeout*/ = TDuration::Seconds(100000)) const override {
        ythrow yexception() << "use WriteLockNode in locks is strange";
    }

    NSaas::TAbstractLock::TPtr ReadLockNode(const TString& /*path*/, TDuration /*timeout*/ = TDuration::Seconds(100000)) const override {
        ythrow yexception() << "use ReadLockNode in locks is strange";
    }

    NSaas::TAbstractLock::TPtr NativeWriteLockNode(const TString& /*path*/, TDuration /*timeout*/ = TDuration::Seconds(100000)) const override {
        ythrow yexception() << "use WriteLockNode in locks is strange";
    }

    NSaas::TAbstractLock::TPtr NativeReadLockNode(const TString& /*path*/, TDuration /*timeout*/ = TDuration::Seconds(100000)) const override {
        ythrow yexception() << "use ReadLockNode in locks is strange";
    }

private:
    NSaas::IVersionedStorage& Slave;
    std::atomic<bool> Disconnected;
};

struct TLockThread : public IThreadFactory::IThreadAble {
    TLockThread(NSaas::IVersionedStorage& storage, bool writeLock)
        : Storage(storage)
        , WriteLock(writeLock)
    {
        Thread = SystemThreadFactory()->Run(this);
    }

    ~TLockThread() {
        Stopped.Signal();
        Thread->Join();
    }

    bool Locked() {
        LockTryed.Wait();
        return !!Lock;
    }

    void DisconnectStorage() {
        Storage.Disconnect();
    }

private:
    void DoExecute() override {
        NSaas::TDistributedLockOptions options;
        Lock = NSaas::TDistributedLock::Create(Storage, "abc", WriteLock, TDuration::Seconds(2), options);
        LockTryed.Signal();
        Stopped.Wait();
    }

    TDisconnectableStorage Storage;
    TManualEvent Stopped;
    TManualEvent LockTryed;
    bool WriteLock;
    TAutoPtr<IThreadFactory::IThread> Thread;
    NSaas::TAbstractLock::TPtr Lock;
};

struct TIncrementer : public IObjectInQueue {
    TIncrementer(NSaas::IVersionedStorage& storage, int& counter)
        : Storage(storage)
        , Counter(counter)
    {}

    void Process(void*) override {
        NSaas::TDistributedLockOptions options;
        auto lock = NSaas::TDistributedLock::Create(Storage, "abc", true, TDuration::Seconds(1000), options);
        UNIT_ASSERT(lock);
        int cntr = Counter;
        ++cntr;
        Sleep(TDuration::MilliSeconds(100));
        Counter = cntr;
    }

    TDisconnectableStorage Storage;
    int& Counter;
};

Y_UNIT_TEST_SUITE(DistributedLockSimple) {

    Y_UNIT_TEST(SingleThread) {
        TDuration lockTimeout = TDuration::Seconds(2);
        NSaas::TStorageOptions options;
        options.LocalOptions.Root = "test";
        Clear("test");
        NSaas::TLocalStorage lStorage(options);
        TDisconnectableStorage storage(lStorage);
        {
            INFO_LOG << "take first lock..." << Endl;
            NSaas::TAbstractLock::TPtr writeLock = NSaas::TDistributedLock::Create(storage, "abc", true, lockTimeout, options.DistrLock);
            INFO_LOG << "take first lock...OK" << Endl;
            UNIT_ASSERT_C(writeLock, "write lock does not asquired");
            INFO_LOG << "take locked locks..." << Endl;
            UNIT_ASSERT_C(!NSaas::TDistributedLock::Create(storage, "abc", true, lockTimeout, options.DistrLock), "two write locks asquired");
            INFO_LOG << "take locked locks...1" << Endl;
            UNIT_ASSERT_C(!NSaas::TDistributedLock::Create(storage, "abc", false, lockTimeout, options.DistrLock), "read lock asquired while write one locked");
            INFO_LOG << "take locked locks...OK" << Endl;
            INFO_LOG << "unlock first lock..." << Endl;
        }
        {
            INFO_LOG << "unlock first lock...OK" << Endl;
            INFO_LOG << "take first read lock..." << Endl;
            NSaas::TAbstractLock::TPtr readLock1 = NSaas::TDistributedLock::Create(storage, "abc", false, lockTimeout, options.DistrLock);
            INFO_LOG << "take first read lock...OK" << Endl;
            UNIT_ASSERT_C(readLock1, "first lock not released");
            INFO_LOG << "take second read lock..." << Endl;
            NSaas::TAbstractLock::TPtr readLock2 = NSaas::TDistributedLock::Create(storage, "abc", false, lockTimeout, options.DistrLock);
            INFO_LOG << "take second read lock...OK" << Endl;
            UNIT_ASSERT_C(readLock2, "read locks prevent each other");
            INFO_LOG << "take locked lock..." << Endl;
            UNIT_ASSERT_C(!NSaas::TDistributedLock::Create(storage, "abc", true, lockTimeout, options.DistrLock), "write lock does asquired while read one locked");
            INFO_LOG << "take locked lock...OK" << Endl;
            INFO_LOG << "unlock read locks..." << Endl;
        }
        INFO_LOG << "unlock read locks...OK" << Endl;
    }

    Y_UNIT_TEST(SimpleMultiThread) {
        NSaas::TStorageOptions options;
        options.LocalOptions.Root = "test";
        Clear("test");
        NSaas::TLocalStorage storage(options);
        {
            INFO_LOG << "take first lock..." << Endl;
            TLockThread writeLock(storage, true);;
            UNIT_ASSERT_C(writeLock.Locked(), "write lock does not asquired");
            INFO_LOG << "take first lock...OK" << Endl;
            INFO_LOG << "take locked locks..." << Endl;
            UNIT_ASSERT_C(!TLockThread(storage, true).Locked(), "two write locks asquired");
            INFO_LOG << "take locked locks...1" << Endl;
            UNIT_ASSERT_C(!TLockThread(storage, false).Locked(), "read lock asquired while write one locked");
            INFO_LOG << "take locked locks...OK" << Endl;
            INFO_LOG << "unlock first lock..." << Endl;
        }
        {
            INFO_LOG << "unlock first lock...OK" << Endl;
            INFO_LOG << "take first read lock..." << Endl;
            TLockThread readLock1(storage, false);
            UNIT_ASSERT_C(readLock1.Locked(), "first lock not released");
            INFO_LOG << "take first read lock...OK" << Endl;
            INFO_LOG << "take second read lock..." << Endl;
            TLockThread readLock2(storage, false);
            UNIT_ASSERT_C(readLock2.Locked(), "read locks prevent each other");
            INFO_LOG << "take second read lock...OK" << Endl;
            INFO_LOG << "take locked lock..." << Endl;
            UNIT_ASSERT_C(!TLockThread(storage, true).Locked(), "write lock does asquired while read one locked");
            INFO_LOG << "take locked lock...OK" << Endl;
            INFO_LOG << "unlock read locks..." << Endl;
        }
        INFO_LOG << "unlock read locks...OK" << Endl;
    }

    Y_UNIT_TEST(Disconnect) {
        TDuration lockTimeout = TDuration::Seconds(2);
        NSaas::TStorageOptions options;
        options.LocalOptions.Root = "test";
        Clear("test");
        NSaas::TLocalStorage lStorage(options);
        TDisconnectableStorage storage(lStorage);
        INFO_LOG << "take first read lock..." << Endl;
        NSaas::TAbstractLock::TPtr readLock1 = NSaas::TDistributedLock::Create(storage, "abc", false, lockTimeout, options.DistrLock);
        UNIT_ASSERT_C(readLock1, "first lock not asqured");
        INFO_LOG << "take first read lock...OK" << Endl;
        TLockThread writeLockThread(lStorage, true);
        Sleep(TDuration::MilliSeconds(500));
        writeLockThread.DisconnectStorage();
        INFO_LOG << "take second read lock..." << Endl;
        NSaas::TAbstractLock::TPtr readLock2 = NSaas::TDistributedLock::Create(storage, "abc", false, TDuration::MilliSeconds(1000), options.DistrLock);
        UNIT_ASSERT_C(readLock2, "second lock not asqured");
        INFO_LOG << "take second read lock...OK" << Endl;
    }

    Y_UNIT_TEST(HardMultithread) {
        NSaas::TStorageOptions options;
        options.LocalOptions.Root = "test";
        Clear("test");
        NSaas::TLocalStorage storage(options);
        int counter = 0;
        TSimpleThreadPool queue;
        queue.Start(0);
        for (int i = 0; i < 100; ++i)
            queue.SafeAddAndOwn(THolder(new TIncrementer(storage, counter)));
        queue.Stop();
        UNIT_ASSERT_EQUAL_C(counter, 100, counter);
    }
}
