#pragma once

#include <cstdint>

#include <memory>

#include <optional>

#include <string>

#include <utility>

#include <vector>

#include "envoy/common/optref.h"

#include "envoy/common/random_generator.h"

#include "envoy/common/time.h"

#include "envoy/config/core/v3/base.pb.h"

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

#include "envoy/http/header_map.h"

#include "envoy/runtime/runtime.h"

#include "envoy/upstream/load_balancer.h"

#include "envoy/upstream/upstream.h"

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

#include "source/common/config/metadata.h"

#include "source/extensions/load_balancing_policies/override_host/override_host_filter_state.h"

#include "absl/status/status.h"

#include "absl/status/statusor.h"

#include "absl/strings/string_view.h"

#include "absl/types/optional.h"

namespace Envoy {

namespace Extensions {

namespace LoadBalancingPolicies {

namespace OverrideHost {

using ::envoy::extensions::load_balancing_policies::override_host::v3::OverrideHost;

using ::Envoy::Random::RandomGenerator;

using ::Envoy::Runtime::Loader;

using ::Envoy::Server::Configuration::ServerFactoryContext;

using ::Envoy::Upstream::ClusterInfo;

using ::Envoy::Upstream::ClusterLbStats;

using ::Envoy::Upstream::Host;

using ::Envoy::Upstream::HostConstSharedPtr;

using ::Envoy::Upstream::HostSelectionResponse;

using ::Envoy::Upstream::LoadBalancerConfigPtr;

using ::Envoy::Upstream::LoadBalancerContext;

using ::Envoy::Upstream::LoadBalancerFactorySharedPtr;

using ::Envoy::Upstream::LoadBalancerParams;

using ::Envoy::Upstream::LoadBalancerPtr;

using ::Envoy::Upstream::PrioritySet;

using ::Envoy::Upstream::ThreadAwareLoadBalancerPtr;

using ::Envoy::Upstream::TypedLoadBalancerFactory;

// Parsed configuration for the dynamic forwarding load balancer. It contains

// factory and config for the load balancer specified in the

// `fallback_policy` field of the OverrideHost config

// proto.

class OverrideHostLbConfig : public Upstream::LoadBalancerConfig {

public:

  struct OverrideSource {

    static OverrideSource make(const OverrideHost::OverrideHostSource& config);

    absl::optional<Http::LowerCaseString> header_name;

    absl::optional<Config::MetadataKey> metadata_key;

  };

  static absl::StatusOr<std::unique_ptr<OverrideHostLbConfig>> make(const OverrideHost& config,

                                                                    ServerFactoryContext& context);

  ThreadAwareLoadBalancerPtr create(const ClusterInfo& cluster_info,

                                    const PrioritySet& priority_set, Loader& runtime,

                                    RandomGenerator& random, TimeSource& time_source) const;

private:

  OverrideHostLbConfig(std::vector<OverrideSource>&& override_host_sources,

                       absl::optional<Config::MetadataKey>&& selected_host_key,

                       TypedLoadBalancerFactory* fallback_load_balancer_factory,

                       LoadBalancerConfigPtr&& fallback_load_balancer_config);

  static absl::StatusOr<std::vector<OverrideSource>> makeOverrideSources(

      const Protobuf::RepeatedPtrField<OverrideHost::OverrideHostSource>& override_sources);

  // Group the factory and config together to make them const in the

  // configuration object.

  struct FallbackLbConfig {

    TypedLoadBalancerFactory* const load_balancer_factory = nullptr;

    const LoadBalancerConfigPtr load_balancer_config;

  };

  const FallbackLbConfig fallback_picker_lb_config_;

  const std::vector<OverrideSource> override_host_sources_;

  const absl::optional<Config::MetadataKey> selected_host_key_;

};

// Load balancer for the dynamic forwarding, supporting external endpoint

// selection by LbTrafficExtension.

// The load balancer uses host list supplied in the request metadata under the

// `com.google.envoy.override_host.localities_and_endpoints` key to pick

// the next backend. TODO(yavlasov): Add a link to the proto describing the

// format of the host list when it is committed.

//

// If the metadata is not present, it falls back to using the load balancer

// specified in the `fallback_policy` field of the

// LoadBalancingPolicyConfig config proto.

// The metadata is not present in two scenarios:

// 1. The Endpoint Picker LbTrafficExtension extension has not been called yet.

//    In this case the picked locality is used to call Endpoint Picker specific

//    to selected locality.

// 2. The Locality Picker extension failed to select a host. In this case, the

//    `fallback_policy` is used as a fallback to pick the backend.

//

// Once the initial locality is picked, the load balancer will use the host list

// from the request metadata to pick the next backend.

class OverrideHostLoadBalancer : public Upstream::ThreadAwareLoadBalancer,

                                 protected Logger::Loggable<Logger::Id::upstream> {

public:

  OverrideHostLoadBalancer(const OverrideHostLbConfig& config,

                           ThreadAwareLoadBalancerPtr fallback_picker_lb)

  LoadBalancerFactorySharedPtr factory() override;

  absl::Status initialize() override;

private:

  // Thread-local LB implementation.

  class LoadBalancerImpl : public Upstream::LoadBalancer {

  public:

    LoadBalancerImpl(const OverrideHostLbConfig& config, LoadBalancerPtr fallback_picker_lb,

                     const PrioritySet& priority_set)

    HostConstSharedPtr peekAnotherHost(LoadBalancerContext* context) override;

    HostSelectionResponse chooseHost(LoadBalancerContext* context) override;

      // TODO(yavlasov): Check if this is needed by the fallback LB.

    absl::optional<Upstream::SelectedPoolAndConnection>

      // This functionality is not supported by dynamic forwarding LB.

  private:

    HostConstSharedPtr getEndpoint(OverrideHostFilterState& override_host_state);

    HostConstSharedPtr findHost(absl::string_view endpoint);

    HostSelectionResponse chooseHostInternal(LoadBalancerContext* context);

    void addSelectedHostKey(LoadBalancerContext* context, HostSelectionResponse& response);

    // Lookup the list of endpoints selected by the LbTrafficExtension in the

    // header or in the request metadata.

    // TODO(wbpcode): will absl::InlinedVector be used here be better?

    std::vector<std::string> getSelectedHosts(LoadBalancerContext* context);

    // Return a list of endpoints selected by the LbTrafficExtension.

    // nullopt if the metadata is not present.

    absl::optional<absl::string_view>

    getSelectedHostsFromMetadata(const ::envoy::config::core::v3::Metadata& metadata,

                                 const Config::MetadataKey& metadata_key);

    // Return a list of endpoints selected by the LbTrafficExtension, specified.

    // in the header. nullopt if the header is not present.

    absl::optional<absl::string_view>

    getSelectedHostsFromHeader(const Http::RequestHeaderMap* header_map,

                               const Http::LowerCaseString& header_name);

    const OverrideHostLbConfig& config_;

    const LoadBalancerPtr fallback_picker_lb_;

    const PrioritySet& priority_set_;

  };

  // LoadBalancerFactory implementation shared by worker threads to create

  // thread-local LB instances. Shared state in this class MUST be protected by

  // appropriate mutexes.

  class LoadBalancerFactoryImpl : public Upstream::LoadBalancerFactory {

  public:

    LoadBalancerFactoryImpl(const OverrideHostLbConfig& config,

                            LoadBalancerFactorySharedPtr fallback_picker_lb_factory)

    // Called by worker threads to create a thread-local load balancer.

    LoadBalancerPtr create(LoadBalancerParams params) override;

  private:

    // Hosts in the load balancer. Owned by the cluster manager.

    const OverrideHostLbConfig& config_;

    LoadBalancerFactorySharedPtr fallback_picker_lb_factory_;

  };

  const OverrideHostLbConfig& config_;

  // Shared factory used to create new thread-local LB implementations.

  std::shared_ptr<LoadBalancerFactoryImpl> factory_;

  const ThreadAwareLoadBalancerPtr fallback_picker_lb_;

};

} // namespace OverrideHost

} // namespace LoadBalancingPolicies

} // namespace Extensions

} // namespace Envoy