#include "expires.h"
#include "set_cookie.h"

#include <balancer/kernel/cookie/set_cookie.h_serialized.h>

#include <balancer/kernel/cookie/utils/utils.h>

#include <util/generic/buffer.h>
#include <util/generic/xrange.h>
#include <util/string/ascii.h>
#include <util/string/cast.h>
#include <util/string/strip.h>

#include <algorithm>
#include <array>

namespace NSrvKernel {
    namespace {
        constexpr TStringBuf APath = "; Path=";
        constexpr TStringBuf ADomain = "; Domain=";
        constexpr TStringBuf AExpires = "; Expires=";
        constexpr TStringBuf AMaxAge = "; Max-Age=";
        constexpr TStringBuf ASecure = "; Secure";
        constexpr TStringBuf AHttpOnly = "; HttpOnly";

        constexpr TStringBuf LoPath = "path";
        constexpr TStringBuf LoDomain = "domain";
        constexpr TStringBuf LoExpires = "expires";
        constexpr TStringBuf LoMaxAge = "max-age";
        constexpr TStringBuf LoSameSite = "samesite";
        constexpr TStringBuf LoSecure = "secure";
        constexpr TStringBuf LoHttpOnly = "httponly";

        constexpr TStringBuf ASameSiteNone = "; SameSite=None";
        constexpr TStringBuf ASameSiteLax = "; SameSite=Lax";
        constexpr TStringBuf ASameSiteStrict = "; SameSite=Strict";

        inline bool IsEqNoCase(TStringBuf name, TStringBuf val) {
            if (name.size() == val.size()) {
                bool res = true;
                for (ui32 i = 0; i < val.size(); ++i) {
                    res &= ((name[i] == val[i]) | (name[i] == (val[i] - 'a' + 'A')));
                }
                return res;
            }
            return false;
        }

        inline bool IsWs(const char* c) noexcept {
            return ui8(*c) <= 0x20;
        }

        [[nodiscard]]
        inline TBlob SubBlobFromValue(TStringBuf value, TBlob data) noexcept {
            if (!value) {
                return {};
            }
            auto valueOff = value.data() - (const char*)data.Data();
            // will crash if value does not belong to data
            return data.SubBlob(valueOff, valueOff + value.size());
        }

        [[nodiscard]]
        inline bool IsSafeAscii(TStringBuf str) noexcept {
            bool isSafe = true;
            for (auto c : str) {
                isSafe &= bool((ui8(c) >= 0x20) & (ui8(c) <= 0x7e));
            }
            return isSafe;
        }

        [[nodiscard]]
        inline bool CheckDomain(TStringBuf domain) noexcept {
            // Can be a domain name or an ipv4/v6 address. We will not try validating this here.
            // Anyway, a syntax error in the Domain attribute will cause all the browsers to ignore the set-cookie.
            // The only case we really need to care about is empty Domain.
            // In this case browsers will accept the cookie but will treat it as a host-only one.
            return !!domain;
        }

        [[nodiscard]]
        inline bool CheckPath(TStringBuf path) noexcept {
            // RFC 6265:
            //      If the attribute-value is empty or if the first character of the
            //      attribute-value is not %x2F ("/"):
            //         Let cookie-path be the default-path.
            //      Otherwise:
            //         Let cookie-path be the attribute-value.
            return path && path[0] == '/';
        }

        [[nodiscard]]
        inline TMaybe<ui32> ParseMaxAgeValue(TStringBuf maxAge) noexcept {
            ui32 res = 0;
            if (!TryFromString(maxAge, res)) {
                return Nothing();
            }
            return res;
        }

        [[nodiscard]]
        inline ESameSite ParseSameSiteValue(TStringBuf sameSite) noexcept {
            if (IsEqNoCase(sameSite, TStringBuf("none"))) {
                return ESameSite::None;
            }
            if (IsEqNoCase(sameSite, TStringBuf("lax"))) {
                return ESameSite::Lax;
            }
            if (IsEqNoCase(sameSite, TStringBuf("strict"))) {
                return ESameSite::Strict;
            }
            // FirstPartyLax and FirstPartyStrict values were drafted a year ago but the draft has since expired
            // and chromium does not support them as of 2020.05.11.
            // Chromium code also mentions an experimental Extended value which is also kept around only for
            // compatibility reasons.
            return ESameSite::Undefined;
        }
    }

    size_t TSetCookie::RenderMaxSz() const noexcept {
        size_t sz = Name.size() + 1 + Value.Size();

        sz += (!Path.empty() ? APath.size() + Path.size() : 0);
        sz += (!Domain.empty() ? ADomain.size() + Domain.size() : 0);

        if (Expires) {
            sz += AExpires.size();
            // Depending on what is allowed in balancer's Expires grammar,
            // SourceInput might as well be longer than the canonical one.
            sz += std::max(PastDate.size(), Expires.Raw().Size());
        }

        sz += (MaxAge ? AMaxAge.size() + 20 : 0);

        constexpr std::array<ui32, GetEnumItemsCount<ESameSite>()> ssLens{
            0, ASameSiteNone.size(), ASameSiteLax.size(), ASameSiteStrict.size()
        };
        static_assert(GetEnumItemsCount<ESameSite>() == 4);
        sz += ssLens[(ui32)SameSite];

        sz += (Secure ? ASecure.size() : 0);
        sz += (HttpOnly ? AHttpOnly.size() : 0);
        return sz;
    }

    TBlob TSetCookie::Render(bool singleThreadedResult) const noexcept {
        const auto sz = RenderMaxSz();
        TBuffer res;
        res.Resize(sz);
        char* pos = res.data();

        auto appendSBuf = [&](TStringBuf val) {
            if (Y_LIKELY(val)) {
                memcpy(pos, val.data(), val.size());
                pos += val.size();
            }
        };
        auto appendChar = [&](char c) {
            *pos++ = c;
        };

        appendSBuf(Name);
        appendChar('=');
        appendSBuf(NCookie::ToStringBuf(Value));

        if (Path.size()) {
            appendSBuf(APath);
            appendSBuf(Path);
        }

        if (Domain.size()) {
            appendSBuf(ADomain);
            appendSBuf(Domain);
        }

        if (Expires) {
            appendSBuf(AExpires);
            if (Expires.Raw().Size()) {
                appendSBuf(NCookie::ToStringBuf(Expires.Raw()));
            } else {
                appendSBuf(NCookie::ToStringBuf(Expires.Get().Render()));
            }
        }

        if (MaxAge) {
            char buf[32] = {}; // 20 actually: -9223372036854775808, 9223372036854775807
            const auto sz = IntToString<10>(*MaxAge, buf, sizeof(buf));
            appendSBuf(AMaxAge);
            appendSBuf({buf, sz});
        }

        switch (SameSite) {
        default:
            break;
        case ESameSite::None:
            appendSBuf(ASameSiteNone);
            break;
        case ESameSite::Lax:
            appendSBuf(ASameSiteLax);
            break;
        case ESameSite::Strict:
            appendSBuf(ASameSiteStrict);
            break;
        }

        if (Secure) {
            appendSBuf(ASecure);
        }

        if (HttpOnly) {
            appendSBuf(AHttpOnly);
        }

        Y_VERIFY(size_t(pos - res.data()) <= sz);
        res.Resize(pos - res.data());
        return singleThreadedResult ? TBlob::FromBufferSingleThreaded(res) : TBlob::FromBuffer(res);
    }

    std::variant<TSetCookie, TSetCookieSyntaxErrors> TSetCookie::Parse(TBlob data) noexcept {
        // TODO(velavokr): first parse in a simple ragel, use this parser as a fallback if the ragel one fails.
        TSetCookieSyntaxErrors errs;
        TSetCookie res;

        TStringBuf setCookie = TStringBuf((const char*)data.Data(), data.Size());

        if (Y_UNLIKELY(!IsSafeAscii(setCookie))) {
            errs.Set(ESetCookieSyntaxError::InvalidOctet);
        }

        setCookie = StripString(setCookie, IsWs);

        TStringBuf attrs;

        {
            TStringBuf nameValue;
            setCookie.Split(';', nameValue, attrs);

            TStringBuf name, value;
            if (!nameValue.TrySplit('=', name, value)) {
                value = nameValue;
            }

            name = StripStringRight(name, IsWs);
            value = StripString(value, IsWs);

            if (!name) {
                if (value) {
                    errs.Set(ESetCookieSyntaxError::CookieNameEmpty);
                } else {
                    errs.Set(ESetCookieSyntaxError::CookieEmpty);
                }
            }

            res.Name.assign(name.data(), name.size());
            res.Value = SubBlobFromValue(value, data);
        }

        bool seenPath = false;
        bool seenDomain = false;
        bool seenExpires = false;
        bool seenMaxAge = false;
        bool seenSameSite = false;

        while (attrs) {
            // Conveniently allows at most one trailing ';'
            TStringBuf attr = attrs.NextTok(';');
            TStringBuf name, value;
            attr.Split('=', name, value);

            name = StripString(name, IsWs);
            value = StripString(value, IsWs);

            if (!name) {
                if (value) {
                    errs.Set(ESetCookieSyntaxError::AttrNameEmpty);
                } else {
                    errs.Set(ESetCookieSyntaxError::AttrEmpty);
                }
                continue;
            }

            // #1 frequent
            if (!seenPath && IsEqNoCase(name, LoPath)) {
                seenPath = true;
                res.Path.assign(value.data(), value.size());
                if (Y_UNLIKELY(!CheckPath(res.Path))) {
                    errs.Set(ESetCookieSyntaxError::PathValueInvalid);
                }
                continue;
            }

            // #2 frequent
            if (!seenDomain && IsEqNoCase(name, LoDomain)) {
                seenDomain = true;
                NCookie::StrToLower(res.Domain, value);
                if (Y_LIKELY(!res.Name.starts_with(CookiePrefixHostOnly)) && Y_UNLIKELY(!CheckDomain(res.Domain))) {
                    errs.Set(ESetCookieSyntaxError::DomainValueInvalid);
                }
                continue;
            }

            // #3 frequent
            if (!seenExpires && IsEqNoCase(name, LoExpires)) {
                seenExpires = true;
                auto raw = SubBlobFromValue(value, data);
                auto expires = TExpires::Parse(raw);
                if (Y_LIKELY(expires)) {
                    res.Expires = TMaybeExpires(raw, *expires);
                } else {
                    errs.Set(ESetCookieSyntaxError::ExpiresValueInvalid);
                }
                continue;
            }

            // #4 frequent
            if (!res.Secure && IsEqNoCase(name, LoSecure)) {
                res.Secure = true;
                continue;
            }

            // #5 frequent
            if (!res.HttpOnly && IsEqNoCase(name, LoHttpOnly)) {
                res.HttpOnly = true;
                continue;
            }

            // #6 frequent
            if (!seenSameSite && IsEqNoCase(name, LoSameSite)) {
                seenSameSite = true;
                res.SameSite = ParseSameSiteValue(value);
                if (Y_UNLIKELY(res.SameSite == ESameSite::Undefined)) {
                    errs.Set(ESetCookieSyntaxError::SameSiteValueInvalid);
                }
                continue;
            }

            // #7 frequent
            if (!seenMaxAge && IsEqNoCase(name, LoMaxAge)) {
                seenMaxAge = true;
                res.MaxAge = ParseMaxAgeValue(value);
                if (Y_UNLIKELY(!res.MaxAge)) {
                    errs.Set(ESetCookieSyntaxError::MaxAgeValueInvalid);
                }
                continue;
            }

            // Should never make it down here on a valid input.

            // We will fail the parsing if there are two conflicting attributes.
            if (IsEqNoCase(name, LoPath)) {
                if (!CheckPath(value)) {
                    errs.Set(ESetCookieSyntaxError::PathValueInvalid);
                }
                if (value != TStringBuf(res.Path)) {
                    errs.Set(ESetCookieSyntaxError::PathValueConflict);
                }
                continue;
            }
            if (IsEqNoCase(name, LoDomain)) {
                if (!CheckDomain(value)) {
                    errs.Set(ESetCookieSyntaxError::DomainValueInvalid);
                }
                if (NCookie::StrToLower(value) != res.Domain) {
                    errs.Set(ESetCookieSyntaxError::DomainValueConflict);
                }
                continue;
            }
            if (IsEqNoCase(name, LoExpires)) {
                auto exp = TExpires::Parse(TBlob::NoCopy(value.Data(), value.size()));
                if (!exp) {
                    errs.Set(ESetCookieSyntaxError::ExpiresValueInvalid);
                }
                if (exp != res.Expires) {
                    errs.Set(ESetCookieSyntaxError::ExpiresValueConflict);
                }
                continue;
            }
            if (IsEqNoCase(name, LoSameSite)) {
                auto ss = ParseSameSiteValue(value);
                if (ss == ESameSite::Undefined) {
                    errs.Set(ESetCookieSyntaxError::SameSiteValueInvalid);
                }
                if (ss != res.SameSite) {
                    errs.Set(ESetCookieSyntaxError::SameSiteValueConflict);
                }
                continue;
            }
            if (IsEqNoCase(name, LoMaxAge)) {
                auto mx = ParseMaxAgeValue(value);
                if (!mx) {
                    errs.Set(ESetCookieSyntaxError::MaxAgeValueInvalid);
                }
                if (mx != res.MaxAge) {
                    errs.Set(ESetCookieSyntaxError::MaxAgeValueConflict);
                }
                continue;
            }
            if (!IsEqNoCase(name, LoSecure) && !IsEqNoCase(name, LoHttpOnly)) {
                // There is also a Priority attribute but its draft expired 4 years ago so perhaps we can safely
                // ignore it as of 2020.05.11. Also it seems the 6265bis spec has been pretty stable for the whole year
                // so far and there are no traces of any recent experimentation in the chromium code so it is unlikely
                // for any new attributes to appear the next few years.
                errs.Set(ESetCookieSyntaxError::AttrUnknown);
            }
        }

        if (!errs.Empty()) {
            return errs;
        }

        return res;
    }

    bool HasDeletion(const TSetCookie& c, TInstant now) noexcept {
        if (c.MaxAge) {
            return *c.MaxAge == 0;
        }
        return c.Expires && c.Expires.Get() < TExpiresNow::Get(now);
    }
}

template <>
void Out<NSrvKernel::TSetCookie>(IOutputStream& out, const NSrvKernel::TSetCookie& c) {
    out << NSrvKernel::NCookie::ToStringBuf(c.Render());
}
