#include <algorithm>

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

#include <util/stream/output.h>
#include <util/string/cast.h>
#include <util/system/condvar.h>
#include <util/system/guard.h>
#include <util/system/mutex.h>
#include <util/thread/factory.h>
#include <util/system/event.h>

#include "zookeeper.h"

namespace NWebmaster {

TZooKeeperClient::TZooKeeperClient(const TString &connectionString, size_t timeout)
    : ConnectionString(connectionString)
    , Timeout(timeout)
{
    NewSession();
}

TZooKeeperClient::~TZooKeeperClient() {
}

void TZooKeeperClient::NewSession() {
    SetState(ST_NEW_SESSION);
    Client.Reset(new NZooKeeper::TZooKeeper(ConnectionString, Timeout, this));
}

TString TZooKeeperClient::Create(const TString &path, const TString &data, NZooKeeper::ECreateMode createMode) {
    EnsureConnected();
    return Client->Create(path, data, NZooKeeper::ACLS_ALL, createMode);
}

void TZooKeeperClient::EnsureConnected() const {
    if (State != ST_CONNECTED) {
        if (!WaitState(ST_CONNECTED, TDuration::MilliSeconds(Timeout))) {
            ythrow yexception() << "not connected";
        }
    }
}

std::pair<bool, int> TZooKeeperClient::Exists(const TString &path, bool watcher) {
    EnsureConnected();

    NZooKeeper::TStatPtr stat;

    try {
        stat = Client->Exists(path, watcher);
    } catch (const std::exception &) {
        throw;
    }

    if (!stat) {
        return std::make_pair(false, 0);
    }

    return std::make_pair(true, stat->version);
}

void TZooKeeperClient::GetChildren(const TString &path, TVector<TString> &children) {
    EnsureConnected();
    children = Client->GetChildren(path, false);
}


void TZooKeeperClient::Delete(const TString &path, int version) {
    EnsureConnected();
    Client->Delete(path, version);
}

void TZooKeeperClient::SetState(EState state) {
    TGuard<TMutex> guard(StateMutex);
    State = state;
    StateCond.BroadCast();
}

bool TZooKeeperClient::WaitState(EState state, TDuration timeOut) const {
    TInstant deadLine = timeOut.ToDeadLine();
    TGuard<TMutex> guard(StateMutex);

    while (State != state) {
        if (!StateCond.WaitD(StateMutex, deadLine)) {
            return false;
        }
    }

    return true;
}

void TZooKeeperClient::Process(const NZooKeeper::TWatchedEvent &event) {
    switch (event.EventType) {
    case NZooKeeper::ET_NONE:
        switch (event.KeeperState) {
        case NZooKeeper::KS_EXPIRED:
            NewSession();
            break;
        case NZooKeeper::KS_NO_SYNC_CONNECTED:
            SetState(ST_RECONNECT);
            break;
        case NZooKeeper::KS_SYNC_CONNECTED:
            SetState(ST_CONNECTED);
            break;
        default:
            ythrow yexception() << "unexpected state " << event.ToString();
        }
        break;
    case NZooKeeper::ET_NOT_WATCHING:
        break;
    case NZooKeeper::ET_NODE_CREATED:
        break;
    case NZooKeeper::ET_NODE_DELETED:
        OnNodeDeleted.Signal();
        break;
    case NZooKeeper::ET_NODE_DATA_CHANGED:
        break;
    case NZooKeeper::ET_NODE_CHILDREN_CHANGED:
        break;
    default:
        ythrow yexception() << "unexpected event " << event.ToString();
    }
}

TZkMutex::TZkMutex(const TString &connectionString, const TString &path, const TString &lockPrefix, size_t timeout)
    : Path(path)
    , LockPrefix(lockPrefix)
    , Timeout(timeout)
    , ZooKeeper(connectionString, Timeout)
{
}

bool TZkMutex::GetSequenceId(const TString &path, int &seq) {
    size_t divider = path.find(LockPrefix);
    if (divider == TString::npos) {
        return false;
    }

    seq = ::FromString(path.substr(divider + LockPrefix.size()));

    return true;
}

bool TZkMutex::GetMinSequenceId(const TVector<TString> &children, int &seq) {
    seq = std::numeric_limits<int>::max();
    bool changed = false;

    for (const TString &child : children) {
        int next;
        if (GetSequenceId(child, next) && next < seq) {
            seq = next;
            changed = true;
        }
    }

    return changed;
}

void TZkMutex::Acquire() {
    if (!LockedPath.empty()) {
        ythrow yexception() << "There is already locked path: " << LockedPath;
    }

    try {
        ZooKeeper.Create(Path, "", NZooKeeper::CM_PERSISTENT);
    } catch (const NZooKeeper::TNodeExistsError &) {
    }

    LockedPath = ZooKeeper.Create(Path + "/" + LockPrefix, "", NZooKeeper::CM_EPHEMERAL_SEQUENTIAL);

    while(true) {
        TVector<TString> children;
        ZooKeeper.GetChildren(Path, children);

        if (std::find(children.begin(), children.end(), LockedPath.substr(Path.size() + 1)) == children.end()) {
            LockedPath = ZooKeeper.Create(Path + "/" + LockPrefix, "", NZooKeeper::CM_EPHEMERAL_SEQUENTIAL);
            continue;
        }

        int mySeq, minSeq;

        if (GetSequenceId(LockedPath, mySeq) && GetMinSequenceId(children, minSeq) && mySeq == minSeq) {
            return; //lock acquired
        }

        if (ZooKeeper.Exists(Path + "/" + children.front(), /*watcher*/ true).first) {
            ZooKeeper.OnNodeDeleted.WaitT(TDuration::MilliSeconds(100));
        }
    }
}

void TZkMutex::Release() {
    try {
        ZooKeeper.Delete(LockedPath, -1);
    } catch(...) {
    }

    LockedPath.clear();
}

bool TZkMutex::Alive() {
    if (LockedPath.empty()) {
        return false;
    }

    return ZooKeeper.Exists(LockedPath, /*watcher*/ false).first;
}

} //namespace NWebmaster
