// Copyright 2023 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef CORE_INTERNAL_ENDPOINT_MANAGER_H_
#define CORE_INTERNAL_ENDPOINT_MANAGER_H_

#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/base/thread_annotations.h"
#include "absl/container/flat_hash_map.h"
#include "absl/functional/any_invocable.h"
#include "absl/time/time.h"
#include "connections/implementation/analytics/packet_meta_data.h"
#include "connections/implementation/client_proxy.h"
#include "connections/implementation/endpoint_channel.h"
#include "connections/implementation/endpoint_channel_manager.h"
#include "connections/implementation/proto/offline_wire_formats.pb.h"
#include "connections/listeners.h"
#include "internal/platform/byte_array.h"
#include "internal/platform/condition_variable.h"
#include "internal/platform/count_down_latch.h"
#include "internal/platform/runnable.h"
#include "internal/platform/single_thread_executor.h"

namespace nearby {
namespace connections {

// Manages all operations related to the remote endpoints with which we are
// interacting.
//
// All processing of incoming and outgoing payloads is spread across this and
// the PayloadManager as described below.
//
// The sending of outgoing payloads originates in
// PayloadManager::SendPayload() before control is transferred over to
// EndpointManager::SendPayloadChunk(). This work happens on one of three
// dedicated writer threads belonging to the PayloadManager. The writer thread
// that is used depends on the PayloadType.
//
// The EndpointManager has one dedicated reader thread for each registered
// endpoint, and the receiving of every incoming payload (and its subsequent
// chunks) originates on one of those threads before control is transferred over
// to PayloadManager::ProcessFrame() (still running on that
// same dedicated reader thread).

class EndpointManager {
 public:
  using OfflineFrame = ::location::nearby::connections::OfflineFrame;

  class FrameProcessor {
   public:
    virtual ~FrameProcessor() = default;

    // @EndpointManagerReaderThread
    // Called for every incoming frame of registered type.
    // NOTE(OfflineFrame& frame):
    // For large payload in data phase, resources may be saved if data is moved,
    // rather than copied (if passing data by reference is not an option).
    // To achieve that, OfflineFrame needs to be either mutabe lvalue reference,
    // or rvalue reference. Rvalue references are discouraged by go/cstyle,
    // and that leaves us with mutable lvalue reference.
    virtual void OnIncomingFrame(
        location::nearby::connections::OfflineFrame& offline_frame,
        const std::string& from_endpoint_id, ClientProxy* to_client,
        location::nearby::proto::connections::Medium current_medium,
        analytics::PacketMetaData& packet_meta_data) = 0;

    // Implementations must call barrier.CountDown() once
    // they're done. This parallelizes the disconnection event across all frame
    // processors.
    //
    // @EndpointManagerThread
    virtual void OnEndpointDisconnect(ClientProxy* client,
                                      const std::string& service_id,
                                      const std::string& endpoint_id,
                                      CountDownLatch barrier,
                                      DisconnectionReason reason) = 0;
  };

  explicit EndpointManager(EndpointChannelManager* manager);
  ~EndpointManager();

  // Invoked from the constructors of the various *Manager components that make
  // up the OfflineServiceController implementation.
  // FrameProcessor* instances are of dynamic duration and survive all sessions.
  // Blocks until registration is complete.
  void RegisterFrameProcessor(
      location::nearby::connections::V1Frame::FrameType frame_type,
      FrameProcessor* processor);
  void UnregisterFrameProcessor(
      location::nearby::connections::V1Frame::FrameType frame_type,
      const FrameProcessor* processor);

  // Invoked from the different PcpHandler implementations (of which there can
  // be only one at a time).
  // Blocks until registration is complete.
  void RegisterEndpoint(ClientProxy* client, const std::string& endpoint_id,
                        const ConnectionResponseInfo& info,
                        const ConnectionOptions& connection_options,
                        std::unique_ptr<EndpointChannel> channel,
                        const ConnectionListener& listener,
                        const std::string& connection_token);
  // Called when a client explicitly asks to disconnect from this endpoint. In
  // this case, we do not notify the client of onDisconnected().
  void UnregisterEndpoint(ClientProxy* client, const std::string& endpoint_id);

  // Returns the maximum supported transmit packet size(MTU) for the underlying
  // transport.
  int GetMaxTransmitPacketSize(const std::string& endpoint_id);

  // Returns the list of endpoints to which sending this chunk failed.
  //
  // Invoked from the PayloadManager's sendPayload() method.
  std::vector<std::string> SendPayloadChunk(
      const location::nearby::connections::PayloadTransferFrame::PayloadHeader&
          payload_header,
      const location::nearby::connections::PayloadTransferFrame::PayloadChunk&
          payload_chunk,
      const std::vector<std::string>& endpoint_ids,
      analytics::PacketMetaData& packet_meta_data);
  std::vector<std::string> SendControlMessage(
      const location::nearby::connections::PayloadTransferFrame::PayloadHeader&
          payload_header,
      const location::nearby::connections::PayloadTransferFrame::ControlMessage&
          control_message,
      const std::vector<std::string>& endpoint_ids);
  // Receiver sends this frame when all the packets are received. Returns the
  // list of endpoints to which sending this frame failed.
  std::vector<std::string> SendPayloadAck(
      std::int64_t payload_id, const std::vector<std::string>& endpoint_ids);
  // Called when we internally want to get rid of the endpoint, without the
  // client directly telling us to. For example...
  //    a) We failed to read from the endpoint in its dedicated reader thread.
  //    b) We failed to write to the endpoint in PayloadManager.
  //    c) The connection was rejected in PCPHandler.
  //    d) The dedicated KeepAlive thread exceeded its period of inactivity.
  // Or in the numerous other cases where a failure occurred and we no longer
  // believe the endpoint is in a healthy state.
  //
  // Note: This must not block. Otherwise we can get into a deadlock where we
  // ask everyone who's registered an FrameProcessor to
  // processEndpointDisconnection() while the caller of DiscardEndpoint() is
  // blocked here.
  void DiscardEndpoint(ClientProxy* client, const std::string& endpoint_id,
                       DisconnectionReason reason);

 protected:
  // For unit tests only to control executing tasks on the executor.
  EndpointManager(EndpointChannelManager* manager,
                  std::unique_ptr<SingleThreadExecutor> serial_executor);

 private:
  class EndpointState {
   public:
    EndpointState(const std::string& endpoint_id,
                  EndpointChannelManager* channel_manager)
        : endpoint_id_{endpoint_id},
          channel_manager_{channel_manager},
          keep_alive_waiter_mutex_{std::make_unique<Mutex>()},
          keep_alive_waiter_{std::make_unique<ConditionVariable>(
              keep_alive_waiter_mutex_.get())} {}

    EndpointState(const EndpointState&) = delete;
    // The default move constructor would not reset |channel_manager_|, for
    // example. This needs to be nullified so the destructor shutdown logic is
    // bypassed when objects are moved.
    EndpointState(EndpointState&& other)
        : endpoint_id_{std::move(other.endpoint_id_)},
          channel_manager_{std::exchange(other.channel_manager_, nullptr)},
          reader_thread_{std::move(other.reader_thread_)},
          keep_alive_waiter_mutex_{
              std::exchange(other.keep_alive_waiter_mutex_, nullptr)},
          keep_alive_waiter_{std::exchange(other.keep_alive_waiter_, nullptr)},
          keep_alive_thread_{std::move(other.keep_alive_thread_)} {}
    EndpointState& operator=(const EndpointState&) = delete;
    EndpointState&& operator=(EndpointState&&) = delete;
    ~EndpointState();

    void StartEndpointReader(Runnable&& runnable);
    void StartEndpointKeepAliveManager(
        absl::AnyInvocable<void(Mutex*, ConditionVariable*)> runnable);

   private:
    const std::string endpoint_id_;
    EndpointChannelManager* channel_manager_;
    SingleThreadExecutor reader_thread_;

    // Use a condition variable so we can wait on the thread but still be able
    // to wake it up before shutting down. We don't want to just sleep and risk
    // blocking shutdown. Note: Create the mutex/condition variable on the heap
    // so raw pointers sent to HandleKeepAlive() aren't invalidated during
    // std::move operations.
    mutable std::unique_ptr<Mutex> keep_alive_waiter_mutex_;
    std::unique_ptr<ConditionVariable> keep_alive_waiter_;
    SingleThreadExecutor keep_alive_thread_;
  };

  // RAII accessor for FrameProcessor
  class LockedFrameProcessor;

  // Provides a mutex per FrameProcessor to prevent unregistering (and
  // destroying) a FrameProcessor when it's in use.
  class FrameProcessorWithMutex {
   public:
    explicit FrameProcessorWithMutex(FrameProcessor* frame_processor = nullptr)
        : frame_processor_{frame_processor} {}

   private:
    FrameProcessor* frame_processor_;
    Mutex mutex_;
    friend class LockedFrameProcessor;
  };

  LockedFrameProcessor GetFrameProcessor(
      location::nearby::connections::V1Frame::FrameType frame_type);

  ExceptionOr<bool> HandleData(const std::string& endpoint_id,
                               ClientProxy* client_proxy,
                               EndpointChannel* endpoint_channel);

  ExceptionOr<bool> HandleKeepAlive(EndpointChannel* endpoint_channel,
                                    absl::Duration keep_alive_interval,
                                    absl::Duration keep_alive_timeout,
                                    Mutex* keep_alive_waiter_mutex,
                                    ConditionVariable* keep_alive_waiter);

  // Waits for a given endpoint EndpointChannelLoopRunnable() workers to
  // terminate.
  // Is called from RegisterEndpoint to avoid races; also called from
  // RemoveEndpoint as part of proper endpoint shutdown sequence.
  // @EndpointManagerThread
  void RemoveEndpointState(const std::string& endpoint_id);

  void EndpointChannelLoopRunnable(
      const std::string& runnable_name, ClientProxy* client_proxy,
      const std::string& endpoint_id,
      absl::AnyInvocable<ExceptionOr<bool>(EndpointChannel*)> handler);

  static void WaitForLatch(const std::string& method_name,
                           CountDownLatch* latch);
  static void WaitForLatch(const std::string& method_name,
                           CountDownLatch* latch, std::int32_t timeout_millis);

  // It should be noted that this method may be called multiple times (because
  // invoking this method closes the endpoint channel, which causes the
  // dedicated reader and KeepAlive threads to terminate, which in turn leads to
  // this method being called), but that's alright because the implementation of
  // this method is idempotent.
  // @EndpointManagerThread
  void RemoveEndpoint(ClientProxy* client, const std::string& endpoint_id,
                      bool notify, DisconnectionReason reason);
  bool ApplySafeToDisconnect(const std::string& endpoint_id,
                             EndpointChannel* endpoint_channel,
                             DisconnectionReason reason);
  void WaitForEndpointDisconnectionProcessing(ClientProxy* client,
                                              const std::string& service_id,
                                              const std::string& endpoint_id,
                                              DisconnectionReason reason);
  void ProcessDisconnectionFrame(
      ClientProxy* client, const std::string& endpoint_id,
      EndpointChannel* endpoint_channel,
      location::nearby::connections::OfflineFrame& frame);
  CountDownLatch NotifyFrameProcessorsOnEndpointDisconnect(
      ClientProxy* client, const std::string& service_id,
      const std::string& endpoint_id, DisconnectionReason reason);

  std::vector<std::string> SendTransferFrameBytes(
      const std::vector<std::string>& endpoint_ids,
      const std::string& payload_transfer_frame_bytes, std::int64_t payload_id,
      std::int64_t offset, const std::string& packet_type,
      analytics::PacketMetaData& packet_meta_data);

  // Executes all jobs sequentially, on a serial_executor_.
  void RunOnEndpointManagerThread(const std::string& name, Runnable runnable);

  ExceptionOr<OfflineFrame> TryDecryptFrame(const ByteArray& data,
                                            EndpointChannel* endpoint_channel);
  EndpointChannelManager* channel_manager_;

  RecursiveMutex frame_processors_lock_;
  absl::flat_hash_map<location::nearby::connections::V1Frame::FrameType,
                      FrameProcessorWithMutex>
      frame_processors_ ABSL_GUARDED_BY(frame_processors_lock_);

  // We keep track of all registered channel endpoints here.
  absl::flat_hash_map<std::string, EndpointState> endpoints_;

  // Indicates whether the destructor has been called yet. If `is_shutdown_`
  // is true, assume any `ClientProxy` pointers are invalid, and should not
  // be used.
  //
  // The ordering of these objects is important: `serial_executor_` must be
  // destroyed before `is_shutdown_` because `serial_executor_` runs all
  // pending tasks during it's destruction, and the "discard-endpoints"
  // task checks `is_shutdown_` to prevent accessing an invalid `ClientProxy`
  // pointer.
  mutable RecursiveMutex mutex_;
  bool is_shutdown_ ABSL_GUARDED_BY(mutex_) = false;

  std::unique_ptr<SingleThreadExecutor> serial_executor_;
};

// Operator overloads when comparing FrameProcessor*.
bool operator==(const EndpointManager::FrameProcessor& lhs,
                const EndpointManager::FrameProcessor& rhs);
bool operator<(const EndpointManager::FrameProcessor& lhs,
               const EndpointManager::FrameProcessor& rhs);

}  // namespace connections
}  // namespace nearby

#endif  // CORE_INTERNAL_ENDPOINT_MANAGER_H_