#include "yt_snapshot_manager.h"

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

#include <mapreduce/yt/common/config.h>
#include <mapreduce/yt/interface/client.h>

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

#include <google/protobuf/text_format.h>

namespace NSaas {

class TYTSnapshotManager : 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<TYTSnapshotManager>;
    static TRegistrator Registrator;
};

TYTSnapshotManager::TRegistrator TYTSnapshotManager::Registrator("yt");

void TYTSnapshotManager::DoPublishSnapshot(const TSnapshot& snapshot) {
    auto client = NYT::CreateClient(Context.Server);
    PublishYtSnapshot(*client, Context.Path, snapshot);
}

void PublishYtSnapshot(NYT::IClientBase& client, const TString& pubDir, const IIndexSnapshotManager::TSnapshot& snapshot) {
    const auto node = TString::Join(pubDir, "/", ToString(snapshot.Timestamp.Seconds()));
    TString data;
    if (!google::protobuf::TextFormat::PrintToString(snapshot.Shards, &data)) {
        ythrow yexception() << "Cannot print TSnapshot protobuf";
    }
    client.Create(node, NYT::NT_FILE, NYT::TCreateOptions().Recursive(true).Force(true));
    auto writer = client.CreateFileWriter(node);
    writer->Write(data.data(), data.size());
}

TYTSnapshotManager::TSnapshots TYTSnapshotManager::DoGetSnapshots(const std::function<bool(TInstant)>& filter, const IStopCondition& stopper) {
    TSnapshots snapshots;
    auto client = NYT::CreateClient(Context.Server, NYT::TCreateClientOptions().Token(Context.Token));
    // TODO(SAAS-5701): global YT config should not be modified by a library
    if (!Context.YTHosts.empty()) {
        NYT::TConfig::Get()->Hosts = Context.YTHosts;
    }
    return GetYtSnapshots(*client, Context.Path, filter, stopper);
}

IIndexSnapshotManager::TSnapshots GetYtSnapshots(NYT::IClientBase& client, const TString& pubDir, const std::function<bool(TInstant)>& filter, const IStopCondition& stopper) {
    IIndexSnapshotManager::TSnapshots snapshots;
    auto nodeList = client.List(pubDir);
    for (const auto& node : nodeList) {
        if (stopper.Stopped()) {
            return {};
        }
        const TString& timestampString = node.AsString();
        const auto ypath = TString::Join(pubDir, "/", timestampString);
        ui64 timestampSeconds;
        if (!TryFromString<ui64>(timestampString, timestampSeconds)) {
            WARNING_LOG << "Skipping " << ypath << ": cannot parse timestamp" << Endl;
        }
        TInstant timestamp = TInstant::Seconds(timestampSeconds);
        if (!filter(timestamp)) {
            continue;
        }
        NRTYServer::TShards shards;
        auto reader = client.CreateFileReader(ypath);
        if (stopper.Stopped()) { // creating a file reader could be quite long
            return {};
        }
        google::protobuf::TextFormat::Parser protoParser;
        protoParser.AllowUnknownField(true);
        if (!protoParser.ParseFromString(reader->ReadAll(), &shards)) {
            INFO_LOG << "Skipping " << ypath << ": not a valid protobuf text" << Endl;
            continue;
        }
        TMap<ui32, IIndexSnapshotManager::TSnapshot> snapshotById;
        TMap<NUtil::TInterval<ui64> , ui32> lastSnapshotIdByShard;
        for (const auto& resource : shards.GetShard()) {
            NUtil::TInterval<ui64> shardInterval(resource.GetShardMin(), resource.GetShardMax());
            NRTYServer::TShardResource* protoCopy = snapshotById[lastSnapshotIdByShard[shardInterval]].Shards.MutableShard()->Add();
            *protoCopy = resource;
            lastSnapshotIdByShard[shardInterval]++;
        }
        for (auto snapshot : snapshotById) {
            snapshot.second.Timestamp = timestamp;
            snapshots.push_back(std::move(snapshot.second));
        }
    }

    return snapshots;
}

static TString ReplacePrefix(TStringBuf src, TStringBuf oldPrefix, TStringBuf newPrefix) {
    Y_ENSURE(src.StartsWith(oldPrefix), "expected '" << src << "' to start with '" << oldPrefix << "'");
    return TString::Join(newPrefix, src.substr(oldPrefix.size()));
}

static void CopyTable(NYT::ICypressClient& cypress, const TString& from, const TString& to) {
    INFO_LOG << "Copying from " << from << " to " << to << Endl;
    cypress.Copy(from, to, NYT::TCopyOptions().Force(true).Recursive(true));
}

IIndexSnapshotManager::TSnapshot CopyDataBlobAndPatchSnapshot(NYT::IClientBase& yt, const IIndexSnapshotManager::TSnapshot& snapshot, const TString& oldDataDir, const TString& newDataDir) {
    INFO_LOG << "Copying " << snapshot.Timestamp.Seconds() << "(" << snapshot.Timestamp << ") data blob from " << newDataDir << " to " << newDataDir << Endl;

    const auto timestampString = ToString(snapshot.Timestamp.Seconds());
    const auto expectedOldBlobPath = TString::Join(oldDataDir, "/", timestampString, ".blob");
    const auto expectedOldBlobResources = TString::Join(oldDataDir, "/", timestampString, ".blob.resources");
    const auto newBlobPath = TString::Join(newDataDir, "/", timestampString, ".blob");
    const auto newBlobResources = TString::Join(newDataDir, "/", timestampString, ".blob.resources");

    Y_ENSURE(yt.Exists(expectedOldBlobPath), expectedOldBlobPath << " YT table not found");

    INFO_LOG << "Patching the copied snapshot, replacing " << expectedOldBlobPath << " with " << newBlobPath << Endl;
    auto newSnapshot = snapshot;
    for (auto& shard : *newSnapshot.Shards.MutableShard()) {
        auto newPath = ReplacePrefix(shard.GetYTIndexTable(), expectedOldBlobPath, newBlobPath);
        shard.SetYTIndexTable(newPath);
    }

    CopyTable(yt, expectedOldBlobPath, newBlobPath);
    if (yt.Exists(expectedOldBlobResources)) {
        CopyTable(yt, expectedOldBlobResources, newBlobResources);
    }

    return newSnapshot;
}

}  // namespace NSaas
