#include "cookie_handler.h"

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

#include <library/cpp/containers/stack_array/stack_array.h>

#include <util/generic/xrange.h>
#include <util/string/builder.h>

#include <alloca.h>

namespace NModCookies {
    using namespace NSrvKernel;
    using namespace NConfig;

    TString TCreateCookie::GenValue(const TConnDescr& descr) const noexcept {
        return std::visit([&](auto&& val) -> TString {
            using TVal = std::decay_t<decltype(val)>;
            if constexpr (std::is_same_v<TVal, TString>) {
                return val;
            } else if constexpr (std::is_same_v<TVal, ECommonFunc>) {
                return ApplyFunc(val, descr);
            }
        }, Value);
    }

    void TCookieHandler::Delete(TString regexp) {
        Delete_ = TFsm::Glue(Delete_, TFsm(regexp, TFsm::TOptions().SetCaseInsensitive(true)));
    }

    void TCookieHandler::Create(TString name, TCreateCookie create, bool weak) {
        const TFsm fsm(name, TFsm::TOptions().SetCaseInsensitive(true));
        if (weak) {
            CreateFuncsWeak_.emplace(name, CreateFuncsWeakOrder_.size());
            CreateFuncsWeakOrder_.emplace_back(name, create);
        } else {
            CreateFuncs_.emplace(name);
            CreateFuncsOrder_.emplace_back(name, create);
        }
    }

    // The cookies module logic is directly mirroring the one of headers module, all quirks included.
    TString TCookieHandler::Apply(TString oldCookie, const TConnDescr& descr) const noexcept {
        TString newCookie;
        newCookie.reserve(oldCookie.size());

        std::vector<bool> excludedCreateFuncOps(CreateFuncs_.size(), false);
        std::vector<bool> excludedCreateWeakFuncOps(CreateFuncsWeak_.size(), false);

        EveryCookieKV([&](TNameValue nv) noexcept {
            bool skipCopyToNew = false;

            // If matched then do not copy or append this cookie
            if (NRegExp::TMatcher(Delete_).Match(nv.Name.GetOrElse(nv.Value)).Final()) {
                skipCopyToNew = true;
            }

            // Skip copying the old value to the new cookie
            // and create new cookie with current name outside this loop.
            if (CreateFuncs_.contains(nv.Name.GetOrElse(nv.Value))) {
                skipCopyToNew = true;
            }

            // Drop all weak cookies.
            // Create weak is not working with delete expression.
            if (auto* fPtr = CreateFuncsWeak_.FindPtr(nv.Name.GetOrElse(nv.Value))) {
                excludedCreateWeakFuncOps[*fPtr] = true;
            }

            if (!skipCopyToNew) {
                AppendCookieKV(newCookie, nv);
            }
        }, oldCookie);

        // The same as strong Create, but only for headers which were not
        // present in source request/response.
        for (size_t idx : xrange(CreateFuncsWeakOrder_.size())) {
            if (excludedCreateWeakFuncOps[idx]) {
                continue;
            }
            auto& op = CreateFuncsWeakOrder_[idx];
            AppendCookieKV(newCookie, {op.first, op.second.GenValue(descr)});
        }

        // Add strong Create cookies. If header already present in cookies
        // it will be added as a separate header.
        for (size_t idx : xrange(CreateFuncsOrder_.size())) {
            if (excludedCreateFuncOps[idx]) {
                continue;
            }
            auto& op = CreateFuncsOrder_[idx];
            AppendCookieKV(newCookie, {op.first, op.second.GenValue(descr)});
        }

        return newCookie;
    }
}
