#include "module.h"
#include "cookie_handler.h"

#include <balancer/kernel/cookie/cookie.h>
#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/http/parser/common_headers.h>
#include <balancer/kernel/http/parser/header_validation.h>
#include <balancer/kernel/http/parser/http.h>
#include <balancer/kernel/module/module.h>

using namespace NConfig;
using namespace NSrvKernel;
using namespace NModCookies;

namespace {
    class TParser: public TModuleParams, public IConfig::IFunc {
    public:
        TParser(TCookieHandler& handler, bool func, bool weak, const TModuleParams& moduleParams)
            : TModuleParams(moduleParams)
            , Handler_(handler)
            , Func_(func)
            , Weak_(weak)
        {
            Config->ForEach(this);
        }

    private:
        START_PARSE {
            Y_ENSURE_EX(key && CheckHeaderName(key),
                TConfigParseError() << "invalid cookie name: " << key.Quote());

            if (Func_) {
                Handler_.Create(key, {::FromString<ECommonFunc>(value->AsString())}, Weak_);
            } else {
                Handler_.Create(key, {value->AsString()}, Weak_);
            }

            return;
        } END_PARSE

    private:
        TCookieHandler& Handler_;
        const bool Func_;
        const bool Weak_;
    };
}

MODULE_BASE(cookies, TModuleWithSubModule) {
    TModule(const TModuleParams& moduleParams)
        : TModuleBase(moduleParams)

    {
        Config->ForEach(this);

        if (!Submodule_) {
            ythrow TConfigParseError() << "no submodule configured";
        }
    }

    START_PARSE {
        if (key == "delete") {
            Handler_.Delete(value->AsString());
            return;
        }

        if (key == "create") {
            TParser(Handler_, false, false, Copy(value->AsSubConfig()));
            return;
        }

        if (key == "create_func") {
            TParser(Handler_, true, false, Copy(value->AsSubConfig()));
            return;
        }

        if (key == "create_weak") {
            TParser(Handler_, false, true, Copy(value->AsSubConfig()));
            return;
        }

        if (key == "create_func_weak") {
            TParser(Handler_, true, true, Copy(value->AsSubConfig()));
            return;
        }

        {
            Submodule_.Reset(Loader->MustLoad(key, Copy(value->AsSubConfig())).Release());
            return;
        }
    } END_PARSE

    TError DoRun(const TConnDescr& descr) const noexcept override {
        auto newDescr = descr.Copy();
        THeaders newHeaders;
        TString allCookies;

        for (const auto& value : newDescr.Request->Headers().GetValuesRef("cookie")) {
            AppendCookieC(allCookies, value.AsStringBuf());
        }

        newDescr.Request->Headers().Replace("cookie", Handler_.Apply(allCookies, descr));

        return Submodule_->Run(newDescr);
    }

private:
    TCookieHandler Handler_;
};

IModuleHandle* NModCookies::Handle() {
    return Ncookies::TModule::Handle();
}
