#include "distributed_lock.h"

#include <library/cpp/digest/md5/md5.h>
#include <library/cpp/logger/global/global.h>

#include <util/generic/hash.h>
#include <util/generic/singleton.h>
#include <util/string/cast.h>
#include <util/string/vector.h>
#include <util/system/mutex.h>
#include <util/thread/pool.h>
#include <util/system/condvar.h>

namespace NSaas {
    namespace {
        struct TCandidate {
            TCandidate(i64 version, bool writeLock, TInstant lastRenew = Now())
                : Version(version)
                , LastRenew(lastRenew)
                , WriteLock(writeLock)
            {}

            TString ToString() const {
                return ::ToString(Version) + "-" + ::ToString(LastRenew.GetValue()) + (WriteLock ? "-write" : "-read");
            }

            static TCandidate Parse(const TString& str) {
                TVector<TString> s = SplitString(str, "-");
                return TCandidate(FromString<i64>(s[0]), s[2] == "write", TInstant::MicroSeconds(FromString<ui64>(s[1])));
            }

            i64 Version;
            TInstant LastRenew;
            bool WriteLock;
        };
    }

    class TDistributedLock::TRenewer {
    Y_DECLARE_SINGLETON_FRIEND();
    public:
        static void Add(TDistributedLock* lock) {
            Singleton<TRenewer>()->AddImpl(lock);
        }

        static void Remove(TDistributedLock* lock) {
            Singleton<TRenewer>()->RemoveImpl(lock);
        }

    private:
        struct TJob : public IObjectInQueue {
            TJob(TRenewer& owner, TDistributedLock* lock)
                : Lock(lock)
                , TimeToProcess(Now())
                , Owner(owner)
            {}

            void Process(void*) override {
                TGuard<TMutex> g(Mutex);
                Stopped.WaitD(Mutex, TimeToProcess);
                if (Lock) {
                    Lock->CheckDeprecated();
                    TInstant start = Now();
                    if (Lock->IsLocked_)
                        Lock->Renew();
                    else
                        Lock->TryLock();
                    TDuration connectDuration = Lock->Config.DisconnectTimeout / Lock->Config.ConnectAttemps;
                    if (Lock->IsLocked_) {
                        TimeToProcess = start + connectDuration;
                        Owner.LockedQueue.SafeAdd(this);
                    } else {
                        TimeToProcess = start + Min<TDuration>(Lock->Config.RetryLockTimeout, connectDuration);
                        Owner.UnlockedQueue.SafeAdd(this);
                    }
                } else {
                    g.Release();
                    THolder<TJob> killer(this);
                }
            }

            void Stop() {
                TGuard<TMutex> g(Mutex);
                CHECK_WITH_LOG(Lock);
                Lock->GiveUp();
                Lock = nullptr;
                Stopped.Signal();
            }

            TDistributedLock* Lock;
            TInstant TimeToProcess;
            TRenewer& Owner;
            TMutex Mutex;
            TCondVar Stopped;
        };

    public:
        ~TRenewer() {
            CHECK_WITH_LOG(Jobs.empty());
            LockedQueue.Stop();
            UnlockedQueue.Stop();
        }

    private:
        TRenewer()
            : LockedQueue("DistrLockLock")
            , UnlockedQueue("DistrLockUnlock")
        {
            LockedQueue.Start(1);
            UnlockedQueue.Start(1);
        }

        void AddImpl(TDistributedLock* lock) {
            TJob* job = new TJob(*this, lock);
            {
                TGuard<TMutex> g(Mutex);
                VERIFY_WITH_LOG(Jobs.insert(std::make_pair(lock, job)).second, "lock added twice");
            }
            job->Process(nullptr);
        }

        void RemoveImpl(TDistributedLock* lock) {
            TGuard<TMutex> g(Mutex);
            auto i = Jobs.find(lock);
            VERIFY_WITH_LOG(i != Jobs.end(), "lock not added");
            TJob* job = i->second;
            Jobs.erase(i);
            g.Release();
            job->Stop();
        }

    private:
        TMutex Mutex;
        TSimpleThreadPool LockedQueue;
        TSimpleThreadPool UnlockedQueue;
        THashMap<TDistributedLock*, TJob*> Jobs;
    };

    TDistributedLock::TDistributedLock(const IVersionedStorage& storage, const TString& id, bool writeLock, const TDistributedLockOptions& config)
        : Storage(storage)
        , Id(id)
        , Path("/locks/distributed/" + MD5::Calc(id))
        , IsLocked_(false)
        , WriteLock(writeLock)
        , Config(config)
    {
        TRenewer::Add(this);
    }

    TAbstractLock::TPtr TDistributedLock::Create(const IVersionedStorage& storage, const TString& id, bool writeLock, const TDuration& timeout, const TDistributedLockOptions& config) {
        TInstant deadline = timeout.ToDeadLine();
        THolder<TDistributedLock> result = THolder<TDistributedLock>(new TDistributedLock(storage, id, writeLock, config));
        return result->WaitLocked(deadline) ? result.Release() : nullptr;
    }

    TDistributedLock::~TDistributedLock() {
        TRenewer::Remove(this);
    }

    bool TDistributedLock::WaitLocked(TInstant deadline) {
        while (!IsLocked_ && Now() < deadline)
            LockedEvent.WaitD(deadline);
        return IsLocked_;
    }

    void TDistributedLock::RemoveCandidate(const TString& name, const char* comment, ui64 line) const {
        if (Storage.RemoveNode(GetCandidatesPath() / name)) {
            DEBUG_LOG << "(" << line << "): " << comment << " candidate " << name << " removed" << Endl;
        } else {
            WARNING_LOG << "(" << line << "): " << comment << " candidate " << name << " not removed" << Endl;
        }
    }

    bool TDistributedLock::CreateCandidate() {
        i64 startSnapshot = 0;
        i64 endSnapshot = 0;
        TString snapshotPath = Path / "snapshot";
        if (!Storage.SetValue(snapshotPath, "", false, false, &startSnapshot)) {
            ERROR_LOG << Id << ": cannot set value to " << snapshotPath << Endl;
            return false;
        }
        TString newLockName;
        if (!Renew(startSnapshot, newLockName))
            return false;
        DEBUG_LOG << "Ephemeral candidate created " << newLockName << Endl;
        if (!Storage.SetValue(snapshotPath, "", false, false, &endSnapshot)) {
            RemoveCandidate(newLockName, "Ephemeral", __LINE__);
            return false;
        }
        TString finalLockName;
        if (!Renew(endSnapshot, finalLockName))
            return false;
        DEBUG_LOG << "Candidate created " << finalLockName << " from ephemeral " << newLockName << Endl;
        RemoveCandidate(newLockName, "Ephemeral", __LINE__);
        LockName = finalLockName;
        return true;
    }

    bool TDistributedLock::Renew(i64 snapshot, TString& newLockName) {
        newLockName = TCandidate(snapshot, WriteLock).ToString();
        TString candidatePath = GetCandidatesPath() / newLockName;
        if (!Storage.SetValue(candidatePath, "", false, false)) {
            ERROR_LOG << Id << ": cannot set value to " << candidatePath << Endl;
            return false;
        }
        return true;
    }

    bool TDistributedLock::Renew() {
        TString newLockName;
        if (!Renew(TCandidate::Parse(LockName).Version, newLockName))
            return false;
        TString oldLockName = LockName;
        LockName = newLockName;
        RemoveCandidate(oldLockName, "Replaced", __LINE__);
        return true;
    }

    void TDistributedLock::TryLock() {
        VERIFY_WITH_LOG(!IsLocked_, "Incorrect usage");
        if (!(LockName ? Renew() : CreateCandidate()))
            return;
        TVector<TString> candidates;
        if (!Storage.GetNodes(GetCandidatesPath(), candidates, true)) {
            ERROR_LOG << Id << ": cannot get children of " << GetCandidatesPath().GetPath() << Endl;
            return;
        }
        TCandidate myCandidate = TCandidate::Parse(LockName);
        ui32 countLocks = 0;
        DEBUG_LOG << "Try lock " << LockName << "..." << Endl;
        for (const TString& candidate : candidates) {
            DEBUG_LOG << "Check candidate " << candidate << Endl;
            if (candidate == LockName)
                continue;
            TCandidate current = TCandidate::Parse(candidate);
            if (Now() - current.LastRenew > Config.DisconnectTimeout) {
                RemoveCandidate(candidate, "Deprecated", __LINE__);
                continue;
            }
            if ((current.Version < myCandidate.Version) && (WriteLock || current.WriteLock))
                ++countLocks;
        }
        DEBUG_LOG << "Try lock " << LockName << "..." << countLocks << Endl;
        IsLocked_ = countLocks == 0;
        if (IsLocked_)
            LockedEvent.Signal();
    }

    void TDistributedLock::GiveUp() {
        if (LockName) {
            RemoveCandidate(LockName, "Released", __LINE__);
            LockName.clear();
        }
    }

    void TDistributedLock::CheckDeprecated() {
        if (!LockName || Now() - TCandidate::Parse(LockName).LastRenew < Config.DisconnectTimeout)
            return;
        VERIFY_WITH_LOG(!IsLocked_, "disconnect from storage with acquired lock. Continue is unsafe");
        GiveUp();
    }

}
