#pragma once

#include <cstdint>

#include <functional>

#include <string>

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

#include "envoy/secret/secret_manager.h"

#include "envoy/server/transport_socket_config.h"

#include "envoy/stats/scope.h"

#include "envoy/thread_local/thread_local.h"

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

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

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

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

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

namespace Envoy {

namespace Upstream {

class OriginalDstClusterFactory;

class OriginalDstClusterTest;

struct HostsForAddress {

  // Primary host for the address. This is set by the first worker that posts

  // to the main to add a host. The field is read by all workers.

  const HostSharedPtr host_;

  // Hosts that are added concurrently with host_ are stored in this list.

  // This is populated by the subsequent workers that have not received the

  // updated table with set host_. The field is only accessed from the main

  // thread.

  std::vector<HostSharedPtr> hosts_;

  // Marks as recently used by load balancers.

  std::atomic<bool> used_;

};

using HostsForAddressSharedPtr = std::shared_ptr<HostsForAddress>;

using HostMultiMap = absl::flat_hash_map<std::string, HostsForAddressSharedPtr>;

using HostMultiMapSharedPtr = std::shared_ptr<HostMultiMap>;

using HostMultiMapConstSharedPtr = std::shared_ptr<const HostMultiMap>;

class OriginalDstCluster;

// Handle object whose sole purpose is to ensure that the destructor of the inner OriginalDstCluster

// is called on the main thread.

class OriginalDstClusterHandle {

public:

  OriginalDstClusterHandle(std::shared_ptr<OriginalDstCluster> cluster)

  ~OriginalDstClusterHandle();

private:

  std::shared_ptr<OriginalDstCluster> cluster_;

  friend class OriginalDstCluster;

};

using OriginalDstClusterHandleSharedPtr = std::shared_ptr<OriginalDstClusterHandle>;

/**

 * The OriginalDstCluster is a dynamic cluster that automatically adds hosts as needed based on the

 * original destination address of the downstream connection. These hosts are also automatically

 * cleaned up after they have not seen traffic for a configurable cleanup interval time

 * ("cleanup_interval_ms").

 */

class OriginalDstCluster : public ClusterImplBase {

public:

  // Upstream::Cluster

  /**

   * Special Load Balancer for Original Dst Cluster.

   *

   * Load balancer gets called with the downstream context which can be used to make sure the

   * Original Dst cluster has a Host for the original destination. Normally load balancers can't

   * modify clusters, but in this case we access a singleton OriginalDstCluster that we can ask to

   * add hosts on demand. Additions are synced with all other threads so that the host set in the

   * cluster remains (eventually) consistent. If multiple threads add a host to the same upstream

   * address then two distinct HostSharedPtr's (with the same upstream IP address) will be added,

   * and both of them will eventually time out.

   */

  class LoadBalancer : public Upstream::LoadBalancer {

  public:

    LoadBalancer(const OriginalDstClusterHandleSharedPtr& parent)

    // Upstream::LoadBalancer

    HostSelectionResponse chooseHost(LoadBalancerContext* context) override;

    // Preconnecting is not implemented for OriginalDstCluster

    // Pool selection not implemented for OriginalDstCluster

    absl::optional<Upstream::SelectedPoolAndConnection>

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

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

    // Lifetime tracking not implemented for OriginalDstCluster

    Network::Address::InstanceConstSharedPtr filterStateOverrideHost(LoadBalancerContext* context);

    Network::Address::InstanceConstSharedPtr requestOverrideHost(LoadBalancerContext* context);

    Network::Address::InstanceConstSharedPtr metadataOverrideHost(LoadBalancerContext* context);

  private:

    const OriginalDstClusterHandleSharedPtr parent_;

    // The optional original host provider that extracts the address from HTTP header map.

    const absl::optional<Http::LowerCaseString>& http_header_name_;

    const absl::optional<Config::MetadataKey>& metadata_key_;

    const absl::optional<uint32_t> port_override_;

    HostMultiMapConstSharedPtr host_map_;

  };

protected:

  OriginalDstCluster(const envoy::config::cluster::v3::Cluster& config,

                     ClusterFactoryContext& context, absl::Status& creation_status);

private:

  friend class OriginalDstClusterFactory;

  friend class OriginalDstClusterTest;

  struct LoadBalancerFactory : public Upstream::LoadBalancerFactory {

    // Upstream::LoadBalancerFactory

    const OriginalDstClusterHandleSharedPtr cluster_;

  };

  struct ThreadAwareLoadBalancer : public Upstream::ThreadAwareLoadBalancer {

    // Upstream::ThreadAwareLoadBalancer

    const OriginalDstClusterHandleSharedPtr cluster_;

  };

  void addHost(HostSharedPtr&);

  void cleanup();

  // ClusterImplBase

  Event::Dispatcher& dispatcher_;

  const std::chrono::milliseconds cleanup_interval_ms_;

  Event::TimerPtr cleanup_timer_;

  absl::Mutex host_map_lock_;

  HostMultiMapConstSharedPtr host_map_ ABSL_GUARDED_BY(host_map_lock_);

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

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

  absl::optional<uint32_t> port_override_;

  friend class OriginalDstClusterFactory;

  friend class OriginalDstClusterHandle;

};

constexpr absl::string_view OriginalDstClusterFilterStateKey =

    "envoy.network.transport_socket.original_dst_address";

class OriginalDstClusterFactory : public ClusterFactoryImplBase {

public:

  absl::StatusOr<std::pair<ClusterImplBaseSharedPtr, ThreadAwareLoadBalancerPtr>>

  createClusterImpl(const envoy::config::cluster::v3::Cluster& cluster,

                    ClusterFactoryContext& context) override;

private:

  friend class OriginalDstClusterTest;

};

DECLARE_FACTORY(OriginalDstClusterFactory);

} // namespace Upstream

} // namespace Envoy