#include "attr_cache.h"

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

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

#include <util/generic/set.h>

namespace NSaas {
    namespace {
        TStringBuf GetParent(TStringBuf path) {
            size_t slashPos = path.rfind('/');
            Y_ENSURE(slashPos != TStringBuf::npos);
            Y_ENSURE(slashPos != 0);
            Y_ENSURE(path.back() != '/');
            return TStringBuf{path.data(), slashPos};
        }

        template <typename TCallback>
        void ReadAttributesImpl(const NYT::TNode& node, const TString& nodePath, TCallback&& callback) {
            if (node.HasAttributes()) {
                callback(nodePath, node.GetAttributes());
            }
            if (node.GetType() == NYT::TNode::Map) {
                for (auto& child : node.AsMap()) {
                    ReadAttributesImpl(child.second, TString::Join(nodePath, '/', child.first), callback);
                }
            }
        }

        template <typename TCallback>
        void ReadAttributes(NYT::ICypressClient& ytClient, const TString& path, const TSet<TString>& attrs, TCallback&& callback) {
            NYT::TGetOptions options;
            options.AttributeFilter_.ConstructInPlace();
            options.AttributeFilter_->Attributes_.assign(attrs.begin(), attrs.end());
            auto node = ytClient.Get(path, options);
            ReadAttributesImpl(node, path, callback);
        }
    }

    struct TAttrCacheNode {
        TMaybe<NYT::TNode> Attributes;
        TInstant ReadTimestamp;

        TAttrCacheNode() = default;
        TAttrCacheNode(TMaybe<NYT::TNode> attrs, TInstant timestamp)
            : Attributes(std::move(attrs))
            , ReadTimestamp(timestamp)
        {
        }

        bool Expired(TDuration lifetime) const {
            return lifetime != TDuration::Zero() && TInstant::Now() - ReadTimestamp > lifetime;
        }
    };

    struct TAttributeCachePrivate {
        TLRUCache<TString, TAttrCacheNode> Cache;
        TSet<TString> Attributes;

        TAttributeCachePrivate(size_t size)
            : Cache(size)
        {
        }
    };

    TAttributeCache::TAttributeCache(size_t size) {
        Y_ENSURE(size > 0);
        P = MakeHolder<TAttributeCachePrivate>(size);
    }

    TAttributeCache::~TAttributeCache() noexcept {
    }

    const NYT::TNode* TAttributeCache::Get(NYT::ICypressClient& ytClient, const TString& path, const TString& attrName, TDuration lifeTime) {
        auto inserted = P->Attributes.insert(attrName);
        if (!inserted.second) {
            auto cached = TryLoadFromCache(path, attrName, lifeTime);
            if (cached.second) {
                return cached.first;
            }
        }
        TInstant timestamp;
        TMaybe<NYT::TNode> result;
        ReadAttributes(ytClient, TString{GetParent(path)}, P->Attributes, [&](const TString& subPath, const NYT::TNode& node) {
            if (timestamp == TInstant{}) {
                timestamp = TInstant::Now();
            }
            if (subPath == path) {
                result = node;
            } else {
                P->Cache.Update(subPath, TAttrCacheNode{node, timestamp});
            }
        });
        P->Cache.Update(path, TAttrCacheNode{std::move(result), timestamp});
        return TryLoadFromCache(path, attrName, lifeTime).first;
    }

    std::pair<const NYT::TNode*, bool> TAttributeCache::TryLoadFromCache(const TString& path, const TString& attrName, TDuration lifeTime) {
        auto it = P->Cache.Find(path);
        if (it == P->Cache.End() || it->Expired(lifeTime)) {
            return {nullptr, false};
        }
        if (!it->Attributes.Defined()) {
            return {nullptr, true};
        }
        auto& attrMap = it->Attributes->AsMap();
        auto attrIt = attrMap.find(attrName);
        auto attr = attrIt == attrMap.end() ? nullptr : &attrIt->second;
        return {attr, true};
    }
}  // namespace NSaas
