#include <util/string/escape.h>
#include "cookie_extractor.h"

namespace NCookie {
    namespace {
        constexpr TStringBuf Domain = "domain";
        constexpr TStringBuf Path = "path";
        constexpr TStringBuf SameSite = "samesite";
        constexpr TStringBuf Secure = "secure";
        constexpr TStringBuf HttpOnly = "httponly";
        constexpr TStringBuf MaxAge = "max-age";
        constexpr TStringBuf Expires = "expires";

        TString MergeSets(TStringBuf a, TStringBuf b) {
            TVector<ui8> res(Reserve(a.size() + b.size()));
            res.insert(res.end(), a.begin(), a.end());
            res.insert(res.end(), b.begin(), b.end());
            SortUnique(res);
            return {(const char*) res.data(), res.size()};
        }

        template <class Get>
        bool SetCookieAttr(TStringBuf token, TStringBuf what, TSetCookie& cookie, Get&& get) {
            if (AsciiHasPrefixIgnoreCase(token, what)) {
                get(cookie)->append(what).append(token.SubStr(what.size()));
                return true;
            }
            return false;
        }

#define Y_SET_COOKIE(token, what, cookie) SetCookieAttr(token, what, cookie, [](auto& c) { return c.Mutable##what(); })
    }

    TCookie ParseCookieToken(TStringBuf token) {
        TCookie cookie;
        cookie.SetInReqCountMax(1);
        cookie.SetReqCount(1);
        cookie.SetTotalCount(1);

        TStringBuf name, value;
        if (token.TrySplit('=', name, value)) {
            TStringBuf len;
            TStringBuf chSet;
            value.Split(':', len, chSet);
            ui32 ln = 0;
            cookie.SetName(name.data(), name.size());
            if (TryFromString<ui32>(len, ln)) {
                cookie.SetValueLenMax(FromString<ui32>(len));
                cookie.SetValueSet(chSet.data(), chSet.size());
            }
        } else {
            cookie.SetName(token.data(), token.size());
            cookie.SetNoValueCount(1);
        }
        return cookie;
    }

    void MergeCookies(TCookie& a, const TCookie& b, bool local) {
        a.SetName(b.GetName());
        a.SetNoValueCount(a.GetNoValueCount() + b.GetNoValueCount());
        a.SetTotalCount(a.GetTotalCount() + b.GetTotalCount());
        if (local) {
            a.SetReqCount(std::max(a.GetReqCount(), b.GetReqCount()));
            a.SetInReqCountMax(a.GetInReqCountMax() + b.GetInReqCountMax());
        } else {
            a.SetReqCount(a.GetReqCount() + b.GetReqCount());
            a.SetInReqCountMax(std::max(a.GetInReqCountMax(), b.GetInReqCountMax()));
        }
        a.SetValueLenMax(std::max(a.GetValueLenMax(), b.GetValueLenMax()));
        a.SetValueSet(MergeSets(a.GetValueSet(), b.GetValueSet()));
    }

    TSetCookie ParseSetCookie(TStringBuf setCookie) {
        TSetCookie res;
        auto nameVal = setCookie.NextTok(';');
        auto cookie = ParseCookieToken(nameVal);
        res.SetName(cookie.GetName());
        res.SetNoValue(cookie.GetNoValueCount());
        res.SetValueLen(cookie.GetValueLenMax());
        res.SetValueSet(cookie.GetValueSet());
        for (auto&& cookieKV : StringSplitter(setCookie).Split(';').SkipEmpty()) {
            TStringBuf token = cookieKV.Token();
            if (Y_SET_COOKIE(token, Domain, res)) {
                continue;
            }
            if (Y_SET_COOKIE(token, Path, res)) {
                continue;
            }
            if (Y_SET_COOKIE(token, SameSite, res)) {
                continue;
            }
            if (Y_SET_COOKIE(token, Secure, res)) {
                continue;
            }
            if (Y_SET_COOKIE(token, HttpOnly, res)) {
                continue;
            }
            if (Y_SET_COOKIE(token, MaxAge, res)) {
                continue;
            }
            if (Y_SET_COOKIE(token, Expires, res)) {
                continue;
            }
            if (res.GetOtherAttrs()) {
                res.MutableOtherAttrs()->append(';');
            }
            res.MutableOtherAttrs()->append(token);
        }
        return res;
    }
}
