#include <saas/library/index_snapshot/base/snapshot_manager.h>

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

class TZooKeeperSnapshotManager : public IIndexSnapshotManager {
    using IIndexSnapshotManager::IIndexSnapshotManager;

protected:
    void DoPublishSnapshot(const TSnapshot& snapshot) override;
    TSnapshots DoGetSnapshots(const std::function<bool(TInstant)>&, const IStopCondition&) override;

private:
    using TRegistrator = IIndexSnapshotManager::TFactory::TRegistrator<TZooKeeperSnapshotManager>;
    static TRegistrator Registrator;
};

TZooKeeperSnapshotManager::TRegistrator TZooKeeperSnapshotManager::Registrator("zookeeper");

void TZooKeeperSnapshotManager::DoPublishSnapshot(const TSnapshot& snapshot) {
    NZooKeeper::TZooKeeper zookeeper(Context.Server);
    TString znode = Context.Path + "/" + ToString(snapshot.Timestamp.Seconds());
    TString data;
    if (!snapshot.Shards.SerializeToString(&data)) {
        ythrow yexception() << "Cannot serialize TSnapshot protobuf";
    }
    if (zookeeper.Exists(znode)) {
        zookeeper.SetData(znode, data);
    } else {
        zookeeper.Create(znode, data, NZooKeeper::ACLS_ALL, NZooKeeper::CM_PERSISTENT);
    }
}

TZooKeeperSnapshotManager::TSnapshots TZooKeeperSnapshotManager::DoGetSnapshots(const std::function<bool(TInstant)>& filter, const IStopCondition& stopper) {
    TSnapshots snapshots;

    NZooKeeper::TZooKeeper zookeeper(Context.Server);
    TVector<TString> snapshotTimestamps = zookeeper.GetChildren(Context.Path);
    for (auto& timestampString : snapshotTimestamps) {
        if (stopper.Stopped()) {
            return {};
        }
        TString znode = Context.Path + "/" + timestampString;
        TInstant timestamp = TInstant::Seconds(FromString<ui64>(timestampString));

        if (!filter(timestamp)) {
            continue;
        }

        const TString& data = zookeeper.GetData(znode);
        TSnapshot snapshot;
        snapshot.Timestamp = timestamp;
        if (!snapshot.Shards.ParseFromString(data)) {
            INFO_LOG << "Skipping " << znode << ": not a valid protobuf" << Endl;
            continue;
        }
        snapshots.push_back(snapshot);
    }

    return snapshots;
}
