#include "by_group_name_from_header.h"

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

namespace NSrvKernel::NByGroupNameFromHeader {

class TPolicy : public IPolicy {
public:
    TPolicy(TString groupName, THolder<IPolicy> slave) noexcept
        : GroupName_(std::move(groupName))
        , Slave_(std::move(slave))
    {}

    void MarkAsRetry() noexcept override {}

    IBackend* Next(IAlgorithm* algo, bool fastAttempt) noexcept override {
        IBackend* backend = algo->NextByName(GroupName_, false);
        if (backend) {
            algo->RemoveSelected(backend);
            return backend;
        }
        return Slave_->Next(algo, fastAttempt);
    }

    void RegisterSuccess() noexcept override {
        Slave_->RegisterSuccess();
    }

    void RegisterFail() noexcept override {
        Slave_->RegisterFail();
    }
private:
    TString GroupName_;
    THolder<IPolicy> Slave_;
};

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

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

    THolder<IPolicy> ConstructPolicy(const TStepParams& params) noexcept override {
        const auto hintValue = params.Descr->Request->Headers().GetFirstValue(BalancingHintHeaderName_);
        // Here we can apply any checks or transformations (i.e. regex matching or rewrite),
        // but for now treat any header value as group name unconditionally.
        // Any transformations of header value can be done in above modules anyway.
        if (hintValue) {
            return MakeHolder<TPolicy>(TString{hintValue}, TPolicyBase::ConstructPolicy(params));
        }

        return TPolicyBase::ConstructPolicy(params);
    }

private:
    START_PARSE {
        ON_KEY("header_name", BalancingHintHeaderName_) {
            return;
        }

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

private:
    TString BalancingHintHeaderName_;
};


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

}  // namespace NSrvKernel::NByGroupNameFromHeader
