#pragma once

#include "envoy/common/callback.h"

#include "envoy/upstream/load_balancer.h"

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

#include "source/extensions/load_balancing_policies/dynamic_modules/lb_config.h"

namespace Envoy {

namespace Extensions {

namespace LoadBalancingPolicies {

namespace DynamicModules {

/**

 * A load balancer implementation that delegates host selection to a dynamic module.

 */

class DynamicModuleLoadBalancer : public Upstream::LoadBalancer,

                                  public Logger::Loggable<Logger::Id::dynamic_modules> {

public:

  DynamicModuleLoadBalancer(DynamicModuleLbConfigSharedPtr config,

                            const Upstream::PrioritySet& priority_set,

                            const std::string& cluster_name);

  ~DynamicModuleLoadBalancer() override;

  // Upstream::LoadBalancer

  Upstream::HostSelectionResponse chooseHost(Upstream::LoadBalancerContext* context) override;

  Upstream::HostConstSharedPtr peekAnotherHost(Upstream::LoadBalancerContext* context) override;

  OptRef<Envoy::Http::ConnectionPool::ConnectionLifetimeCallbacks> lifetimeCallbacks() override;

  absl::optional<Upstream::SelectedPoolAndConnection>

  selectExistingConnection(Upstream::LoadBalancerContext* context, const Upstream::Host& host,

                           std::vector<uint8_t>& hash_key) override;

  // Accessors for callbacks.

  // Per-host custom data storage.

  bool setHostData(uint32_t priority, size_t index, uintptr_t data);

  bool getHostData(uint32_t priority, size_t index, uintptr_t* data) const;

  // Accessors for hosts added/removed during the on_host_membership_update callback.

private:

  DynamicModuleLbConfigSharedPtr config_;

  const Upstream::PrioritySet& priority_set_;

  std::string cluster_name_;

  envoy_dynamic_module_type_lb_module_ptr in_module_lb_;

  // Handle for the member update callback registration. Automatically unregisters on destruction.

  Envoy::Common::CallbackHandlePtr member_update_cb_;

  // Temporary pointers to host vectors, valid only during on_host_membership_update callback.

  const Upstream::HostVector* hosts_added_{};

  const Upstream::HostVector* hosts_removed_{};

  // Per-host data storage keyed by (priority, index). This is per-LB-instance (per-worker).

  absl::flat_hash_map<std::pair<uint32_t, size_t>, uintptr_t> per_host_data_;

};

} // namespace DynamicModules

} // namespace LoadBalancingPolicies

} // namespace Extensions

} // namespace Envoy