#include "lock_unique.h"

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

#include <mapreduce/yt/interface/client.h>
#include <mapreduce/yt/interface/init.h>
#include <mapreduce/yt/client/client.h>

#include <util/generic/yexception.h>

TYtLockUnique::TYtLockUnique(const TString& serverName, const TString& path, TDuration timeout, ui32 attemptCount, const NYT::TCreateClientOptions& options)
        : Locked(false)
        , Path(path)
        , Timeout(timeout)
        , AttemptCount(attemptCount)
{
    Y_ASSERT(!path.empty() && !serverName.empty());

    Client = NYT::CreateClient(serverName, options);

    TString way = TString{TStringBuf(Path).RBefore('/')};

    CreateFolder(way);
}

bool TYtLockUnique::TryLock() {
    for (ui32 attemptNumber = 0; attemptNumber < AttemptCount; ++attemptNumber) {
        if (TryLockOnce()) {
            return true;
        }
    }
    return false;
}

bool TYtLockUnique::TryLockOnce() {
    if (Locked) {
        return true;
    }

    TString way = TString{TStringBuf(Path).RBefore('/')};
    TString childName = TString{TStringBuf(Path).RAfter('/')};

    NYT::TLockOptions lockOptions = NYT::TLockOptions();
    lockOptions.Waitable(true);
    lockOptions.ChildKey(childName);

    NYT::ILockPtr lock;
    try {
        Transaction = Client->StartTransaction();
        lock = Transaction->Lock(way, NYT::LM_SHARED, lockOptions);
    } catch (...) {
        ERROR_LOG << CurrentExceptionMessage() << Endl;
        return false;
    }
    INFO_LOG << "TYtLockUnique " << Path << " : lock created" << Endl;

    try {
        lock->Wait(Timeout);
        INFO_LOG << "TYtLockUnique " << Path << " : Lock is aquired" << Endl;
        Locked = true;
    } catch (...) {
        AbortTransaction();
        INFO_LOG << "TYtLockUnique " << Path << " : Lock wasn't aquired" << Endl;
    }
    return Locked;
}

void TYtLockUnique::Unlock() {
    if (Locked) {
        AbortTransaction();
        INFO_LOG << "TYtLockUnique " << Path << " : Lock is released" << Endl;
        Locked = false;
    }
}

void TYtLockUnique::AbortTransaction() {
    try {
        Transaction->Abort();
    } catch (...) {
        ERROR_LOG << CurrentExceptionMessage() << Endl;
    }
}

TYtLockUnique::~TYtLockUnique() {
    if (Locked) {
        AbortTransaction();
        INFO_LOG << "TYtLockUnique " << Path << " : Lock is released" << Endl;
    }
}

void TYtLockUnique::CreateFolder(const TString& way) {
    NYT::TCreateOptions createOptions;
    createOptions.IgnoreExisting(true);
    createOptions.Recursive(true);
    Client->Create(way, NYT::NT_MAP, createOptions);
}
