#pragma once

#include <algorithm>

#include <bitset>

#include <cstdint>

#include <functional>

#include <map>

#include <memory>

#include <string>

#include "envoy/common/optref.h"

#include "envoy/config/cluster/v3/cluster.pb.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/runtime/runtime.h"

#include "envoy/stats/scope.h"

#include "envoy/stream_info/stream_info.h"

#include "envoy/upstream/load_balancer.h"

#include "source/common/common/macros.h"

#include "source/common/protobuf/protobuf.h"

#include "source/common/protobuf/utility.h"

#include "source/common/upstream/load_balancer_context_base.h"

#include "source/common/upstream/upstream_impl.h"

#include "source/extensions/load_balancing_policies/subset/subset_lb_config.h"

#include "absl/container/node_hash_map.h"

#include "absl/types/optional.h"

namespace Envoy {

namespace Upstream {

using HostHashSet = absl::flat_hash_set<HostSharedPtr>;

class SubsetLoadBalancer : public LoadBalancer, Logger::Loggable<Logger::Id::upstream> {

public:

  SubsetLoadBalancer(const SubsetLoadBalancerConfig& lb_config,

                     const Upstream::ClusterInfo& cluster_info, const PrioritySet& priority_set,

                     const PrioritySet* local_priority_set, ClusterLbStats& stats,

                     Stats::Scope& scope, Runtime::Loader& runtime, Random::RandomGenerator& random,

                     TimeSource& time_source);

  ~SubsetLoadBalancer() override;

  // Upstream::LoadBalancer

  HostSelectionResponse chooseHost(LoadBalancerContext* context) override;

  // TODO(alyssawilk) implement for non-metadata match.

  // Pool selection not implemented.

  absl::optional<Upstream::SelectedPoolAndConnection>

  selectExistingConnection(Upstream::LoadBalancerContext* /*context*/,

                           const Upstream::Host& /*host*/,

  // Lifetime tracking not implemented.

  using SubsetMetadata = std::vector<std::pair<std::string, Protobuf::Value>>;

  static std::string describeMetadata(const SubsetMetadata& kvs);

private:

  struct SubsetSelectorFallbackParams;

  void initSubsetAnyOnce();

  void initSubsetDefaultOnce();

  void initSubsetSelectorMap();

  void initSelectorFallbackSubset(const envoy::config::cluster::v3::Cluster::LbSubsetConfig::

                                      LbSubsetSelector::LbSubsetSelectorFallbackPolicy&);

  HostConstSharedPtr

  chooseHostForSelectorFallbackPolicy(const SubsetSelectorFallbackParams& fallback_params,

                                      LoadBalancerContext* context);

  HostConstSharedPtr chooseHostIteration(LoadBalancerContext* context);

  // Represents a subset of an original HostSet.

  class HostSubsetImpl : public HostSetImpl {

  public:

    HostSubsetImpl(const HostSet& original_host_set, bool locality_weight_aware,

                   bool scale_locality_weight)

    void update(const HostHashSet& matching_hosts, const HostVector& hosts_added,

                const HostVector& hosts_removed);

    LocalityWeightsConstSharedPtr

    determineLocalityWeights(const HostsPerLocality& hosts_per_locality) const;

  private:

    const HostSet& original_host_set_;

    const bool locality_weight_aware_;

    const bool scale_locality_weight_;

  };

  // Represents a subset of an original PrioritySet.

  class PrioritySubsetImpl : public PrioritySetImpl {

  public:

    PrioritySubsetImpl(const SubsetLoadBalancer& subset_lb, bool locality_weight_aware,

                       bool scale_locality_weight);

    void update(uint32_t priority, const HostHashSet& matching_hosts, const HostVector& hosts_added,

                const HostVector& hosts_removed);

    void updateSubset(uint32_t priority, const HostHashSet& matching_hosts,

    // Thread aware LB if applicable.

    ThreadAwareLoadBalancerPtr thread_aware_lb_;

    // Current active LB.

    LoadBalancerPtr lb_;

  protected:

    HostSetImplPtr createHostSet(uint32_t priority, absl::optional<bool> weighted_priority_health,

                                 absl::optional<uint32_t> overprovisioning_factor) override;

  private:

    const PrioritySet& original_priority_set_;

    const PrioritySet* original_local_priority_set_{};

    const bool locality_weight_aware_;

    const bool scale_locality_weight_;

    bool empty_ = true;

  };

  using HostSubsetImplPtr = std::unique_ptr<HostSubsetImpl>;

  using PrioritySubsetImplPtr = std::unique_ptr<PrioritySubsetImpl>;

  class LbSubsetEntry;

  struct SubsetSelectorMap;

  using LbSubsetEntryPtr = std::shared_ptr<LbSubsetEntry>;

  using SubsetSelectorMapPtr = std::shared_ptr<SubsetSelectorMap>;

  using ValueSubsetMap = absl::node_hash_map<HashedValue, LbSubsetEntryPtr>;

  using LbSubsetMap = absl::node_hash_map<std::string, ValueSubsetMap>;

  using SubsetSelectorFallbackParamsRef = std::reference_wrapper<SubsetSelectorFallbackParams>;

  using MetadataFallbacks = Protobuf::RepeatedPtrField<Protobuf::Value>;

public:

  class LoadBalancerContextWrapper : public LoadBalancerContext {

  public:

    LoadBalancerContextWrapper(LoadBalancerContext* wrapped,

                               const std::set<std::string>& filtered_metadata_match_criteria_names);

    LoadBalancerContextWrapper(LoadBalancerContext* wrapped,

                               Router::MetadataMatchCriteriaConstPtr metadata_match_criteria)

    LoadBalancerContextWrapper(LoadBalancerContext* wrapped,

                               const Protobuf::Struct& metadata_match_criteria_override);

    // LoadBalancerContext

    const HealthyAndDegradedLoad& determinePriorityLoad(

        const PrioritySet& priority_set, const HealthyAndDegradedLoad& original_priority_load,

  private:

    LoadBalancerContext* wrapped_;

    Router::MetadataMatchCriteriaConstPtr metadata_match_;

  };

private:

  struct SubsetSelectorFallbackParams {

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

        LbSubsetSelectorFallbackPolicy fallback_policy_;

    const std::set<std::string>* fallback_keys_subset_ = nullptr;

  };

  struct SubsetSelectorMap {

    absl::node_hash_map<std::string, SubsetSelectorMapPtr> subset_keys_;

    SubsetSelectorFallbackParams fallback_params_;

  };

  class LbSubset {

  public:

    virtual HostSelectionResponse chooseHost(LoadBalancerContext* context) const PURE;

    virtual void pushHost(uint32_t priority, HostSharedPtr host) PURE;

    virtual void finalize(uint32_t priority) PURE;

    virtual bool active() const PURE;

  };

  using LbSubsetPtr = std::unique_ptr<LbSubset>;

  class PriorityLbSubset : public LbSubset {

  public:

    PriorityLbSubset(const SubsetLoadBalancer& subset_lb, bool locality_weight_aware,

                     bool scale_locality_weight)

    // Subset

    // Called after pushHost. Update subset by the hosts that pushed in the pushHost. If no any host

    // is pushed then subset_ will be set to empty.

    void finalize(uint32_t priority) override;

    std::vector<std::pair<HostHashSet, HostHashSet>> host_sets_;

    PrioritySubsetImpl subset_;

  };

  class SingleHostLbSubset : public LbSubset {

    // Subset

    // This is called at most once for every update for single host subset.

    // Called after pushHost. Update subset by the host that pushed in the pushHost. If no any host

    // is pushed then subset_ will be set to nullptr.

        // No any host for current subset and priority. Try remove record in the hosts_.

        // Single host is set for current subset and priority.

    // We will update subsets for every priority separately and these simple map can help us

    // to ensure which priority has valid host quickly.

    std::map<uint32_t, HostSharedPtr> hosts_;

    std::map<uint32_t, HostSharedPtr> new_hosts_;

    HostConstSharedPtr subset_;

  };

  // Entry in the subset hierarchy.

  class LbSubsetEntry {

  public:

    LbSubsetMap children_;

    // Only initialized if a match exists at this level.

    LbSubsetPtr lb_subset_;

    // Used to quick check if entry is single host subset entry or not.

    bool single_host_subset_{};

  };

  void initLbSubsetEntryOnce(LbSubsetEntryPtr& entry, bool single_host_subset);

  // Create filtered default subset (if necessary) and other subsets based on current hosts.

  void refreshSubsets();

  void refreshSubsets(uint32_t priority);

  // Called by HostSet::MemberUpdateCb

  void update(uint32_t priority, const HostVector& all_hosts);

  void updateFallbackSubset(uint32_t priority, const HostVector& all_hosts);

  void processSubsets(uint32_t priority, const HostVector& all_hosts);

  HostConstSharedPtr tryChooseHostFromContext(LoadBalancerContext* context, bool& host_chosen);

  absl::optional<SubsetSelectorFallbackParamsRef>

  tryFindSelectorFallbackParams(LoadBalancerContext* context);

  bool hostMatches(const SubsetMetadata& kvs, const Host& host);

  LbSubsetEntryPtr

  findSubset(const std::vector<Router::MetadataMatchCriterionConstSharedPtr>& matches);

  LbSubsetEntryPtr findOrCreateLbSubsetEntry(LbSubsetMap& subsets, const SubsetMetadata& kvs,

                                             uint32_t idx);

  void forEachSubset(LbSubsetMap& subsets, std::function<void(LbSubsetEntryPtr&)> cb);

  void purgeEmptySubsets(LbSubsetMap& subsets);

  std::vector<SubsetMetadata> extractSubsetMetadata(const std::set<std::string>& subset_keys,

                                                    const Host& host);

  HostConstSharedPtr chooseHostWithMetadataFallbacks(LoadBalancerContext* context,

                                                     const MetadataFallbacks& metadata_fallbacks);

  const Protobuf::Value* getMetadataFallbackList(LoadBalancerContext* context) const;

  LoadBalancerContextWrapper removeMetadataFallbackList(LoadBalancerContext* context);

  const SubsetLoadBalancerConfig& lb_config_;

  const Upstream::ClusterInfo& cluster_info_;

  ClusterLbStats& stats_;

  Stats::Scope& scope_;

  Runtime::Loader& runtime_;

  Random::RandomGenerator& random_;

  TimeSource& time_source_;

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

      fallback_policy_;

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

      metadata_fallback_policy_;

  const SubsetMetadata default_subset_metadata_;

  std::vector<SubsetSelectorPtr> subset_selectors_;

  const PrioritySet& original_priority_set_;

  const PrioritySet* original_local_priority_set_;

  Common::CallbackHandlePtr original_priority_set_callback_handle_;

  LbSubsetEntryPtr subset_any_;

  LbSubsetEntryPtr subset_default_;

  // Reference to sub_set_any_ or subset_default_.

  LbSubsetEntryPtr fallback_subset_;

  LbSubsetEntryPtr panic_mode_subset_;

  // Forms a trie-like structure. Requires lexically sorted Host and Route metadata.

  LbSubsetMap subsets_;

  // Forms a trie-like structure of lexically sorted keys+fallback policy from subset

  // selectors configuration

  SubsetSelectorMapPtr selectors_;

  Stats::Gauge* single_duplicate_stat_{};

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

  const bool locality_weight_aware_ : 1;

  const bool scale_locality_weight_ : 1;

  const bool list_as_any_ : 1;

  const bool allow_redundant_keys_{};

};

} // namespace Upstream

} // namespace Envoy