#pragma once

#include <balancer/kernel/client_hints/impl/client_hints.h>
#include <balancer/production/x/collect_cookie_stats/lib/cookie.pb.h>

#include <library/cpp/yson/node/node.h>
#include <library/cpp/yson/node/node_io.h>

#include <util/generic/hash_set.h>
#include <util/string/split.h>
#include <util/string/strip.h>

#include <contrib/libs/re2/re2/stringpiece.h>
#include <contrib/libs/re2/re2/re2.h>

namespace NCookie {

    TCookie ParseCookieToken(TStringBuf token);

    void MergeCookies(TCookie& a, const TCookie& b, bool local);

    TSetCookie ParseSetCookie(TStringBuf setCookie);

    struct THostPath {
        TString Host;
        TString Path;
    };


    template <class OnCookie>
    void IterateCookies(const THashMap<TString, TCookie>& cookies, OnCookie&& onCookie) {
        TVector<TStringBuf> cookieKeys(Reserve(cookies.size()));
        for (const auto& cookie : cookies) {
            cookieKeys.emplace_back(cookie.first);
        }
        Sort(cookieKeys);
        for (auto&& cookie : cookieKeys) {
            onCookie(cookies.find(cookie)->second);
        }
    }

    class THostPathExtractor {
    public:
        THostPath Apply(const NYT::TNode& row) const {
            const auto& host = StripString(row["host"].AsString(), EqualsStripAdapter('"'));
            const auto& query = StripString(row["query"].AsString(), EqualsStripAdapter('"'));
            re2::StringPiece path;
            re2::RE2::FullMatch({query.data(), query.size()}, ExtractPath_, &path);
            return {host, TString(path.data(), path.size())};
        }

    private:
        re2::RE2 ExtractPath_{"[A-Z]+ (/[^?# ]*)(?:[?# ].*)?"};
    };

    class TCookieExtractor {
    public:
        template <class OnCookie>
        void Apply(const NYT::TNode& row, OnCookie&& onCookie) const {
            const auto& workflowCol = row["workflow"].AsString();

            re2::StringPiece workflow{workflowCol.data(), workflowCol.size()};
            re2::StringPiece cookie;
            while (re2::RE2::FindAndConsume(&workflow, ExtractCookie_, &cookie)) {
                for (auto&& cookieKV : StringSplitter(TStringBuf(cookie.data(), cookie.size())).Split(';').SkipEmpty()) {
                    onCookie(ParseCookieToken(cookieKV.Token()));
                }
            }
        }

    private:
        re2::RE2 ExtractCookie_{" <::CookieMeta:(.*?)::>"};
    };


    class TSetCookieExtractor {
    public:
        template <class OnCookie, class OnSetCookie>
        THostPath Apply(const NYT::TNode& row, OnCookie&& onCookie, OnSetCookie&& onSetCookie) const {
            THostPath hostPath = ExtractHostPath_.Apply(row);

            const auto& workflowCol = row["workflow"].AsString();

            THashMap<TString, ui32> cookies;
            ExtractCookie_.Apply(row, [&](const TCookie& c) {
                cookies[c.GetName()] += c.GetTotalCount();
                onCookie(c);
            });

            re2::StringPiece workflow{workflowCol.data(), workflowCol.size()};
            re2::StringPiece userAgent;
            re2::RE2::FindAndConsume(&workflow, ExtractUserAgent_, &userAgent);
            NSrvKernel::TUATraits uaTraits = NSrvKernel::GetUATraits({userAgent.data(), userAgent.size()});
            re2::StringPiece setCookie;
            while (re2::RE2::FindAndConsume(&workflow, ExtractSetCookie_, &setCookie)) {
                auto res = ParseSetCookie({setCookie.data(), setCookie.size()});
                res.SetNeedsSameSite(uaTraits.SameSiteNoneSupport);
                res.SetReqHost(hostPath.Host);
                res.SetReqPath(hostPath.Path);
                res.SetInReqCount(cookies[res.GetName()]);
                onSetCookie(res);
            }

            return hostPath;
        }

    private:
        re2::RE2 ExtractUserAgent_{" <::user-agent:(.*?)::>"};
        re2::RE2 ExtractSetCookie_{" <::SetCookieMeta:(.*?)::>"};
        THostPathExtractor ExtractHostPath_;
        TCookieExtractor ExtractCookie_;
    };
}
