#include "by_name_from_header.h"
#include "by_name.h"

#include <balancer/kernel/balancer/algorithm.h>
#include <balancer/kernel/balancing/backend.h>
#include <balancer/kernel/helpers/common_parsers.h>
#include <balancer/kernel/http/parser/http.h>

#include <balancer/modules/balancer/policies/policy_factory.h>

namespace NSrvKernel::NByNameFromHeader {

POLICY_FACTORY_BASE(by_name_from_header_policy, TPolicyFactoryWithSlave), public TModuleParams, private NByName::TByNamePolicyFactoryBase {
    TPolicyFactory(const TModuleParams& mp)
        : TModuleParams(mp)
    {
        mp.Config->ForEach(this);
        CheckConfiguration();

        if (HintToBackend_.empty()) {
            ythrow TConfigParseError{} << "balancer2/by_name_from_header_policy: no hints specified";
        }

        if (!BalancingHintHeaderName_) {
            ythrow TConfigParseError{} << "balancer2/by_name_from_header_policy: empty header name";
        }

        BalancingHintFsm_ = MakeHolder<NRegExp::TFsm>(BalancingHintHeaderName_, NRegExp::TFsm::TOptions().SetCaseInsensitive(true));
    }

    THolder<IPolicy> ConstructPolicy(const TStepParams& params) noexcept override {
        if (const auto hintValue = params.Descr->Request->Headers().GetFirstValue(*BalancingHintFsm_)) {
            for (const auto& it : HintToBackend_) {
                if (hintValue == it.first) {
                    return MakeHolder<NByName::TByNamePolicy>(it.second, TPolicyBase::ConstructPolicy(params), Opts_);
                }
            }
        }

        return TPolicyBase::ConstructPolicy(params);
    }

private:
    START_PARSE {
        if (key == "hints") {
            ParseMap(value->AsSubConfig(), [&m = HintToBackend_](const auto&, auto* value) {
                TString hint;
                TString backend;

                ParseMap(value->AsSubConfig(), [&hint, &backend](const auto& key, auto* value) {
                    if (key == "hint") {
                        hint = value->AsString();
                    } else if (key == "backend") {
                        backend = value->AsString();
                    } else {
                        ythrow TConfigParseError{} << "unknown key `" << key << "`";
                    }
                });

                if (!hint || !backend) {
                    ythrow TConfigParseError{} << "no `hint` or `backend` found";
                }

                if (!m.insert(std::make_pair(hint, backend)).second) {
                    ythrow TConfigParseError{} << "repeated hint `" << hint << '`';
                }
            });
            return;
        }

        ON_KEY("header_name", BalancingHintHeaderName_) {
            return;
        }

        if (ConsumeOpts(key, value)) {
            return;
        }

        if (Configure(key, Copy(value->AsSubConfig()))) {
            return;
        }
    } END_PARSE

private:
    TMap<TString, TString> HintToBackend_;
    TString BalancingHintHeaderName_ = "X-Yandex-Balancing-Hint";
    THolder<NRegExp::TFsm> BalancingHintFsm_;
};


INodeHandle<IPolicyFactory>* Handle() {
    return TPolicyFactory::Handle();
}

}  // namespace NSrvKernel::NByNameFromHeader
