#pragma once

#include "envoy/config/cluster/v3/cluster.pb.h"

#include "envoy/extensions/load_balancing_policies/common/v3/common.pb.h"

#include "envoy/extensions/load_balancing_policies/common/v3/common.pb.validate.h"

#include "envoy/extensions/load_balancing_policies/subset/v3/subset.pb.h"

#include "envoy/extensions/load_balancing_policies/subset/v3/subset.pb.validate.h"

#include "envoy/upstream/load_balancer.h"

namespace Envoy {

namespace Upstream {

using SubsetLbConfigProto = envoy::extensions::load_balancing_policies::subset::v3::Subset;

using ClusterProto = envoy::config::cluster::v3::Cluster;

using LegacySubsetLbConfigProto = envoy::config::cluster::v3::Cluster::LbSubsetConfig;

/**

 * This is part of subset load balancer config and is used to store config of single selector.

 */

class SubsetSelector {

public:

  SubsetSelector(const Protobuf::RepeatedPtrField<std::string>& selector_keys,

                 envoy::config::cluster::v3::Cluster::LbSubsetConfig::LbSubsetSelector::

                     LbSubsetSelectorFallbackPolicy fallback_policy,

                 const Protobuf::RepeatedPtrField<std::string>& fallback_keys_subset,

                 bool single_host_per_subset);

  // SubsetSelector

  envoy::config::cluster::v3::Cluster::LbSubsetConfig::LbSubsetSelector::

      LbSubsetSelectorFallbackPolicy

private:

  const std::set<std::string> selector_keys_;

  const std::set<std::string> fallback_keys_subset_;

  // Keep small members (bools and enums) at the end of class, to reduce alignment overhead.

  const envoy::config::cluster::v3::Cluster::LbSubsetConfig::LbSubsetSelector::

      LbSubsetSelectorFallbackPolicy fallback_policy_;

  const bool single_host_per_subset_ : 1;

};

using SubsetSelectorPtr = std::shared_ptr<SubsetSelector>;

/**

 * Load Balancer subset configuration.

 */

class LoadBalancerSubsetInfo {

public:

  /**

   * @return bool true if load balancer subsets are configured.

   */

  virtual bool isEnabled() const PURE;

  /**

   * @return LbSubsetFallbackPolicy the fallback policy used when

   * route metadata does not match any subset.

   */

  virtual envoy::config::cluster::v3::Cluster::LbSubsetConfig::LbSubsetFallbackPolicy

  fallbackPolicy() const PURE;

  /**

   * @return LbSubsetMetadataFallbackPolicy the fallback policy used to try different route metadata

   * until a host is found

   */

  virtual envoy::config::cluster::v3::Cluster::LbSubsetConfig::LbSubsetMetadataFallbackPolicy

  metadataFallbackPolicy() const PURE;

  /**

   * @return Protobuf::Struct the struct describing the metadata for a

   *         host to be included in the default subset.

   */

  virtual const Protobuf::Struct& defaultSubset() const PURE;

  /*

   * @return const std:vector<std:set<std::string>>& a vector of

   * sorted keys used to define load balancer subsets.

   */

  virtual const std::vector<SubsetSelectorPtr>& subsetSelectors() const PURE;

  /*

   * @return bool whether routing to subsets should take locality weights into account.

   */

  virtual bool localityWeightAware() const PURE;

  /*

   * @return bool whether the locality weights should be scaled to compensate for the

   * fraction of hosts removed from the original host set.

   */

  virtual bool scaleLocalityWeight() const PURE;

  /*

   * @return bool whether to attempt to select a host from the entire cluster if host

   * selection from the fallback subset fails.

   */

  virtual bool panicModeAny() const PURE;

  /*

   * @return bool whether matching metadata should attempt to match against any of the

   * elements in a list value defined in endpoint metadata.

   */

  virtual bool listAsAny() const PURE;

  /*

   * @return bool whether redundant key/value pairs is allowed in the request metadata.

   */

  virtual bool allowRedundantKeys() const PURE;

};

using LoadBalancerSubsetInfoPtr = std::unique_ptr<LoadBalancerSubsetInfo>;

/**

 * Both the legacy and extension subset proto configs are converted to this class.

 */

class LoadBalancerSubsetInfoImpl : public LoadBalancerSubsetInfo {

public:

  using FallbackPolicy =

      envoy::config::cluster::v3::Cluster::LbSubsetConfig::LbSubsetFallbackPolicy;

  using MetadataFallbackPolicy =

      envoy::config::cluster::v3::Cluster::LbSubsetConfig::LbSubsetMetadataFallbackPolicy;

  using SubsetFallbackPolicy = envoy::config::cluster::v3::Cluster::LbSubsetConfig::

      LbSubsetSelector::LbSubsetSelectorFallbackPolicy;

  LoadBalancerSubsetInfoImpl(const SubsetLbConfigProto& subset_config)

        metadata_fallback_policy_(

      // Sort subset selectors by number of keys, descending. This will ensure that the longest

      // matching subset selector will be at the beginning of the list.

  LoadBalancerSubsetInfoImpl(const LegacySubsetLbConfigProto& subset_config)

  LoadBalancerSubsetInfoImpl()

      : LoadBalancerSubsetInfoImpl(

  // Upstream::LoadBalancerSubsetInfo

private:

  const Protobuf::Struct default_subset_;

  std::vector<SubsetSelectorPtr> subset_selectors_;

  // Keep small members (bools and enums) at the end of class, to reduce alignment overhead.

  const FallbackPolicy fallback_policy_;

  const MetadataFallbackPolicy metadata_fallback_policy_;

  const bool enabled_ : 1;

  const bool locality_weight_aware_ : 1;

  const bool scale_locality_weight_ : 1;

  const bool panic_mode_any_ : 1;

  const bool list_as_any_ : 1;

  const bool allow_redundant_keys_{};

};

using DefaultLoadBalancerSubsetInfo = ConstSingleton<LoadBalancerSubsetInfoImpl>;

class SubsetLoadBalancerConfig : public Upstream::LoadBalancerConfig {

public:

  SubsetLoadBalancerConfig(Server::Configuration::ServerFactoryContext& factory_context,

                           const SubsetLbConfigProto& config, absl::Status& creation_status);

  SubsetLoadBalancerConfig(Server::Configuration::ServerFactoryContext& factory_context,

                           const ClusterProto& cluster, absl::Status& creation_status);

  SubsetLoadBalancerConfig(LoadBalancerSubsetInfoPtr subset_info,

                           TypedLoadBalancerFactory* child_factory,

                           LoadBalancerConfigPtr child_config);

  Upstream::ThreadAwareLoadBalancerPtr

  createLoadBalancer(const Upstream::ClusterInfo& cluster_info,

                     const Upstream::PrioritySet& child_priority_set, Runtime::Loader& runtime,

private:

  LoadBalancerSubsetInfoPtr subset_info_;

  Upstream::TypedLoadBalancerFactory* child_lb_factory_{};

  Upstream::LoadBalancerConfigPtr child_lb_config_;

};

} // namespace Upstream

} // namespace Envoy