#include "module.h"

#include <balancer/kernel/custom_io/limitio.h>
#include <balancer/kernel/helpers/default_instance.h>
#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/http/parser/common_headers.h>
#include <balancer/kernel/http/parser/http.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/regexp/regexp_re2.h>

#include <library/cpp/string_utils/old_url_normalize/url.h>

namespace NSrvKernel {
    namespace {
        struct TXWwwFormUrlencodedFsm : public TFsm, public TWithDefaultInstance<TXWwwFormUrlencodedFsm> {
            TXWwwFormUrlencodedFsm() noexcept
                : TFsm("application/x-www-form-urlencoded", TFsm::TOptions().SetCaseInsensitive(true).SetSurround(true))
            {}
        };

        struct TKeyFsm : public TFsm, public TWithDefaultInstance<TKeyFsm> {
            TKeyFsm() noexcept
                : TFsm("[-0-9a-zA-Z_]+")
            {}
        };

        static TString Checked(TString s, const char* itemName) {
            Y_ENSURE_EX(TMatcher(TKeyFsm::Instance()).Match(s).Final(),
                        TConfigParseError() << "invalid " << itemName << " name " << s.Quote());
            return s;
        }

        class TCgiFromHdr : public TIntrusiveListItem<TCgiFromHdr> {
        public:
            TCgiFromHdr(const TString& cgi, const TString& hdr)
                : Cgi_(Checked(cgi, "cgi"))
                , HdrFsm_(hdr, TFsm::TOptions().SetCaseInsensitive(true))
            {}

            void Apply(TString& cgi, TStringBuf value) const noexcept {
                TChunkPtr encodedValue = NewChunkReserve();

                encodedValue->Shrink(
                    encodedValue->Offset(
                        EncodeRFC1738(
                            encodedValue->Data(), encodedValue->Length(),
                            value.Data(), value.Size()
                        )
                    ) - 1
                );

                if (cgi.Empty()) {
                    cgi += "?";
                } else {
                    cgi += "&";
                }

                cgi += Cgi_;
                cgi += "=";
                cgi += encodedValue->AsStringBuf();
            }

            const TFsm& HdrFsm() const noexcept {
                return HdrFsm_;
            }

        private:
            const TString Cgi_;
            const TFsm HdrFsm_;
        };


        class THdrFromCgi : public TIntrusiveListItem<THdrFromCgi> {
        public:
            THdrFromCgi(const TString& hdr, const TString& cgi)
                : Hdr_(Checked(hdr, "header"))
                , HdrFsm_(Hdr_, TFsm::TOptions().SetCaseInsensitive(true))
                , CgiReg_(
                    TString("(?:^|[?&])").append(Checked(cgi, "cgi")).append("=([^&]*)(?:&|$)"),
                    TRegexp::TOpts().SetCaseInsensitive(true).SetPosixSyntax(false)
                )
            {}

            void Apply(TStringBuf cgi, TRequest& request) const noexcept {
                TVector<TStringBuf> value;

                if (CgiReg_.Extract(cgi, &value, true) && value.size() == 2) {
                    request.Headers().Delete(HdrFsm_);
                    request.Headers().Add(Hdr_, TString(value[1]));
                }
            }

        private:
            const TString Hdr_;
            const TFsm HdrFsm_;
            const TRegexp CgiReg_;
        };


        template <class TItem>
        class THdrCgiList : public IConfig::IFunc {
            START_PARSE
                {
                    Items.PushBack(new TItem(key, value->AsString()));
                    return;
                }
            END_PARSE

        public:
            TIntrusiveListWithAutoDelete<TItem, TDelete> Items;
        };


        using THdrFromCgiList = THdrCgiList<THdrFromCgi>;
        using TCgiFromHdrList = THdrCgiList<TCgiFromHdr>;


        bool IsXWwwFormUrlencoded(const TRequest& req) noexcept {
            if (auto hdrValue = req.Headers().GetFirstValue(TContentTypeFsm::Instance())) {
                if (Match(TXWwwFormUrlencodedFsm::Instance(), hdrValue)) {
                    return true;
                }
            }
            return false;
        }
    }
}

using namespace NConfig;
using namespace NSrvKernel;

MODULE_BASE(hdrcgi, TModuleWithSubModule) {
public:
    TModule(const TModuleParams& mp)
        : TModuleBase(mp)
    {
        TXWwwFormUrlencodedFsm::Instance();
        Config->ForEach(this);

        Y_UNUSED(TKeyFsm::Instance());
    }

private:
    START_PARSE {
        if (key == "cgi_from_hdr") {
            value->AsSubConfig()->ForEach(&CgiFromHdrList_);
            return;
        }

        if (key == "hdr_from_cgi") {
            value->AsSubConfig()->ForEach(&HdrFromCgiList_);
            return;
        }

        ON_KEY("body_scan_limit", BodyScanLimit_) {
            return;
        }

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

    TError DoRun(const TConnDescr& descr) const noexcept override {
        TRequest* const request = descr.Request;
        TString newCgi;

        for (const auto& item : CgiFromHdrList_.Items) {
            if (auto hdrValue = request->Headers().GetFirstValue(item.HdrFsm())) {
                if (newCgi.Empty() && !request->RequestLine().CGI.Empty()) {
                    newCgi = request->RequestLine().CGI.AsStringBuf();
                }
                item.Apply(newCgi, hdrValue);
            }
        }

        for (const auto& item : HdrFromCgiList_.Items) {
            item.Apply(request->RequestLine().CGI.AsStringBuf(), *request);
        }

        if (!newCgi.Empty()) {
            request->RequestLine().CGI = TStringStorage(std::move(newCgi));
        }

        if (BodyScanLimit_ && HdrFromCgiList_.Items && IsXWwwFormUrlencoded(*request) && !request->Props().UpgradeRequested) {
            if (TError error = descr.Input->FillBuffer(BodyScanLimit_).ReleaseError()) {
                descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), "client read error");
                return error;
            }
            TChunkList postBody = descr.Input->RecvBuffered();
            TChunkList toScan = CutPrefix(BodyScanLimit_, postBody);
            {
                for (const auto& item : HdrFromCgiList_.Items) {
                    item.Apply(StrInplace(toScan), *request);
                }
            }
            postBody.Prepend(std::move(toScan));
            descr.Input->UnRecv(std::move(postBody));
        }

        return Submodule_->Run(descr);
    }

private:
    TCgiFromHdrList CgiFromHdrList_;
    THdrFromCgiList HdrFromCgiList_;

    size_t BodyScanLimit_ = 0;
};

IModuleHandle* NModHdrCgi::Handle() {
    return TModule::Handle();
}
