#pragma once

#include <balancer/kernel/balancer/policy.h>

#include <balancer/kernel/module/node.h>
#include <balancer/kernel/ctl/ctl.h>

namespace NSrvKernel {

    const INodeFactory<IPolicyFactory>* PolicyFactories();

    template <class I, class T, const char* N>
    class TPolicyFactoryWithSlave : public TGenericNode<I, T, N> {
    public:
        using TPolicyBase = TPolicyFactoryWithSlave;
        using TGenericNode<I, T, N>::TGenericNode;
        bool Configure(const TString& key, const TModuleParams& mp) {
            THolder<IPolicyFactory> policyFactory = PolicyFactories()->Load(key, mp);
            if (policyFactory) {
                if (SlavePolicyFactory_) {
                    ythrow TConfigParseError{} << "balancer2/" << N << ": trying to replace existing policy with new one (" << key << ')';
                }
                SlavePolicyFactory_.Reset(policyFactory.Release());
                return true;
            }
            return false;
        }
        void CheckConfiguration() {
            if (!SlavePolicyFactory_) {
                ythrow TConfigParseError{} << "balancer2/" << N << ": no subpolicy configured";
            }
        }

        THolder<IPolicy> ConstructPolicy(const TStepParams& params) noexcept override {
            return SlavePolicyFactory_->ConstructPolicy(params);
        }
        void FillFeatures(TPolicyFeatures& features) const noexcept override {
            return SlavePolicyFactory_->FillFeatures(features);
        }
        void Init(IWorkerCtl* process) override {
            SlavePolicyFactory_->Init(process);
        }
    private:
        THolder<IPolicyFactory> SlavePolicyFactory_;
    };

    template <class Tls, class T, const char* N>
    class TPolicyFactoryWithTls : public TPolicyFactoryWithSlave<IPolicyFactory, T, N> {
    public:
        using TPolicyBase = TPolicyFactoryWithTls;

        TPolicyFactoryWithTls() = delete;
        explicit TPolicyFactoryWithTls(size_t workersCount)
            : Tls_(workersCount + 1)
        {}

        void Init(IWorkerCtl* process) noexcept final {
            TPolicyFactoryWithSlave<IPolicyFactory, T, N>::Init(process);
            Y_VERIFY(!Tls_[process->WorkerId()].Get());
            Tls_[process->WorkerId()] = std::move(InitTls(process));
        }

    protected:
        Tls& GetTls(IWorkerCtl* process) const noexcept {
            Y_VERIFY(Tls_[process->WorkerId()].Get());
            return *Tls_[process->WorkerId()];
        }

    private:
        virtual THolder<Tls> InitTls(IWorkerCtl*) = 0;

    private:
        TVector<THolder<Tls>> Tls_;
    };

}  // namespace NSrvKernel


#define POLICY_FACTORY(X) NODEIMPL(IPolicyFactory, TPolicyFactory, X)

#define POLICY_FACTORY_BASE(X, B) NODEIMPL_BASE(IPolicyFactory, TPolicyFactory, X, B)
#define POLICY_FACTORY_WITH_TLS(X, TLS) NODEIMPL_BASE(TLS, TPolicyFactory, X, TPolicyFactoryWithTls)
