#include "module.h"

#include <balancer/kernel/coro/coro_async.h>
#include <balancer/kernel/coro/coro_event.h>
#include <balancer/kernel/custom_io/queue.h>
#include <balancer/kernel/custom_io/rewind.h>
#include <balancer/kernel/helpers/errors.h>
#include <balancer/kernel/http/parser/http.h>
#include <balancer/kernel/http/parser/httpdecoder.h>
#include <balancer/kernel/matcher/matcher.h>
#include <balancer/kernel/module/module.h>
#include <util/generic/ptr.h>
#include <util/generic/vector.h>

using namespace NConfig;
using namespace NSrvKernel;

MODULE(response_matcher) {

    using TPriority = float;
    class TItem {
    public:
        TItem(const TString& name, TPriority priority, THolder<IResponseMatcher> matcher, THolder<IModule> slave)
            : Name_(name)
            , Priority_(priority)
            , Matcher_(std::move(matcher))
            , Slave_(std::move(slave))
        {
            if (!Slave_) {
                ythrow TConfigParseError() << "no module configured";
            }

            if (!Matcher_) {
                ythrow TConfigParseError() << "no matcher configured";
            }

            if (Priority_ < 0.0) {
                ythrow TConfigParseError() << "negative priority";
            }
        }

        const TString& Name() const noexcept {
            return Name_;
        }

        TPriority Priority() const noexcept {
            return Priority_;
        }

        IResponseMatcher* Matcher() const noexcept {
            return Matcher_.Get();
        }

        IModule* Slave() const noexcept {
            return Slave_.Get();
        }

    private:
        TString Name_;
        TPriority Priority_ = 0.0;
        THolder<IResponseMatcher> Matcher_;
        THolder<IModule> Slave_;
    };

    class TItemConfigParser final : public TModuleParams, public IConfig::IFunc {
    public:

        TItemConfigParser(const TString& name, const TModuleParams& mp)
            : TModuleParams(mp)
            , Name_(name)
        {
            Config->ForEach(this);
        }

        TItem Construct() {
            Y_VERIFY(!Done_, "bug in response matcher code");
            Done_ = true;
            return TItem(Name_, Priority_, std::move(Matcher_), std::move(Slave_));
        }
    private:

        START_PARSE {
            ON_KEY("priority", Priority_) {
                return;
            }

            if (auto matcher = ConstructResponseMatcher(nullptr, key, value->AsSubConfig())) {
                Y_ENSURE_EX(!Matcher_, TConfigParseError() << "duplicate matchers in response matcher:" << Name_);
                Matcher_ = std::move(matcher);
                return;
            }

            {
                Y_ENSURE_EX(!Slave_, TConfigParseError() << "duplicate modules in response matcher:" << Name_);
                Loader->MustLoad(key, Copy(value->AsSubConfig())).Swap(Slave_);
                return;
            }
        } END_PARSE

        TString Name_;
        TPriority Priority_ = 0.0;
        THolder<IResponseMatcher> Matcher_;
        THolder<IModule> Slave_;
        bool Done_ = false;
    };

public:
    TModule(const TModuleParams& mp)
        : TModuleBase(mp)
    {
        Config->ForEach(this);

        if (Items_.empty()) {
            ythrow TConfigParseError() << "no on_response configured";
        }
        if (!Submodule_) {
            ythrow TConfigParseError() << "no module configured";
        }
        if (ForwardHeaders_) {
            ForwardHeadersFsm_ = MakeHolder<TFsm>(ForwardHeaders_, TFsm::TOptions().SetCaseInsensitive(true));
        }
        StableSortBy(Items_, [](const TItem& item) { return item.Priority() ; });
    }

    START_PARSE {
        if (key == "module") {
            TSubLoader(Copy(value->AsSubConfig())).Swap(Submodule_);
            return;
        }
        if (key == "on_response") {
            ParseMap(value->AsSubConfig(), [this](const auto& key, auto* value) {
                TItemConfigParser parser(key, Copy(value->AsSubConfig()));
                Items_.emplace_back(parser.Construct());
            });
            return;
        }
        ON_KEY("forward_headers", ForwardHeaders_) {
            return;
        }
        ON_KEY("buffer_size", BufferSize_) {
            return;
        }
    } END_PARSE

    void ForwardHeaders(TResponse& response, TRequest& request) const {
        if (!ForwardHeadersFsm_) {
            return;
        }
        for (auto& header : response.Headers()) {
            if (::Match(*ForwardHeadersFsm_, header.first.AsStringBuf())) {
                request.Headers().Replace(header.first.AsString(), THeaders::MakeOwned(std::move(header.second)));
            }
        }
    }

    TError DoRun(const TConnDescr& descr) const noexcept override {
        if (MayHaveBody(*descr.Request)) {
            return Submodule_->Run(descr);
        }

        TResponse response;
        const TItem* matchedItem = nullptr;

        auto output = MakeHttpOutput([&](TResponse&& cameResponse, const bool cameForceClose, TInstant) {
            if (matchedItem = FindFirstMatch(cameResponse)) {
                response = std::move(cameResponse);
                return TError{};
            }

            return descr.Output->SendHead(std::move(cameResponse), cameForceClose, TInstant::Max());
        }, [&](TChunkList lst, TInstant deadline) {
            if (matchedItem) {
                return TError{};
            }
            return descr.Output->Send(std::move(lst), deadline);
        }, [&](THeaders&& headers, TInstant deadline) -> TError {
            if (matchedItem) {
                return TError{};
            }
            return descr.Output->SendTrailers(std::move(headers), deadline);
        });

        Y_PROPAGATE_ERROR(Submodule_->Run(descr.CopyOut(output)));


        if (matchedItem) {
            ForwardHeaders(response, *descr.Request);
            Y_PROPAGATE_ERROR(matchedItem->Slave()->Run(descr));
        }

        return TError{};
    }

    bool DoExtraAccessLog() const noexcept override {
        return false;
    }

    bool MayHaveBody(const TRequest& request) const noexcept {
        return request.RequestLine().Method != EMethod::GET
             || request.Props().ContentLength > 0
             || request.Props().ChunkedTransfer;
    }

private:
    const TItem* FindFirstMatch(const TResponse& response) const noexcept {
        for (const TItem& item: Items_) {
            if (item.Matcher()->Match(response)) {
                return &item;
            }
        }
        return nullptr;
    }

private:
    ui64 BufferSize_ = 300 * 1024; // 300 Kb
    THolder<IModule> Submodule_;
    TVector<TItem> Items_;
    THolder<TFsm> ForwardHeadersFsm_;
    TString ForwardHeaders_;
};

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