// Copyright 2022 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef COMPONENTS_MIRRORING_SERVICE_OPENSCREEN_SESSION_HOST_H_
#define COMPONENTS_MIRRORING_SERVICE_OPENSCREEN_SESSION_HOST_H_

#include <optional>
#include <vector>

#include "base/component_export.h"
#include "base/functional/callback_forward.h"
#include "base/gtest_prod_util.h"
#include "base/logging.h"
#include "base/memory/weak_ptr.h"
#include "base/sequence_checker.h"
#include "base/task/single_thread_task_runner.h"
#include "base/unguessable_token.h"
#include "components/mirroring/mojom/cast_message_channel.mojom.h"
#include "components/mirroring/mojom/resource_provider.mojom.h"
#include "components/mirroring/mojom/session_observer.mojom.h"
#include "components/mirroring/mojom/session_parameters.mojom.h"
#include "components/mirroring/service/audio_capturing_callback.h"
#include "components/mirroring/service/media_remoter.h"
#include "components/mirroring/service/mirror_settings.h"
#include "components/mirroring/service/mirroring_gpu_factories_factory.h"
#include "components/mirroring/service/mirroring_logger.h"
#include "components/mirroring/service/openscreen_message_port.h"
#include "components/mirroring/service/openscreen_stats_client.h"
#include "components/mirroring/service/rpc_dispatcher.h"
#include "components/openscreen_platform/event_trace_logging_platform.h"
#include "components/openscreen_platform/task_runner.h"
#include "gpu/config/gpu_info.h"
#include "media/capture/video/video_capture_feedback.h"
#include "media/cast/cast_environment.h"
#include "media/cast/constants.h"
#include "media/mojo/mojom/video_encode_accelerator.mojom.h"
#include "mojo/public/cpp/bindings/pending_receiver.h"
#include "mojo/public/cpp/bindings/pending_remote.h"
#include "mojo/public/cpp/bindings/remote.h"
#include "services/network/public/mojom/socket_factory.mojom.h"
#include "third_party/openscreen/src/cast/streaming/public/sender_session.h"

using openscreen::cast::capture_recommendations::Recommendations;

namespace base {
class OneShotTimer;
}

namespace media {
class AudioInputDevice;
namespace cast {
class AudioSender;
class VideoSender;
}  // namespace cast
}  // namespace media

namespace viz {
class Gpu;
}  // namespace viz

namespace mirroring {

class RpcDispatcher;
class VideoCaptureClient;

// Minimum required bitrate used for calculating bandwidth.
inline constexpr int kMinRequiredBitrate = 384 << 10;  // 384 kbps

// Default bitrate used before we have a calculation.
inline constexpr int kDefaultBitrate = 5000000;  // 5 mbps

// Hosts a streaming session by hosting an `openscreen::cast::SenderSession` and
// doing all of the necessary interfacing for audio and video capture, switching
// between mirroring and remoting, and setting up audio and video streams to
// encode and send captured content.
//
// On construction, an Open Screen SenderSession is immediately created and
// negotiation of a streaming session is started. The session host will stay
// in a good state until either the mirroring service notices a disconnection
// and tears down this streaming session, or a fatal error occurs.
//
// NOTE: most methods should be called on the same sequence as construction.
// This class also uses additional task runners, such as the IO task runner of
// this utility process for accessing the GPU, and dedicated video and audio
// encoder threads. Finally, some methods such as
// AudioCapturingCallback::Capture may be called on the audio thread.
class COMPONENT_EXPORT(MIRRORING_SERVICE) OpenscreenSessionHost final
    : public openscreen::cast::SenderSession::Client,
      public MediaRemoter::Client {
 public:
  // NOTE: some notes on constructor arguments:
  //    `session_params`: connection information for the receiver.
  //    `max_resolution`: width and height that should never be exceeded.
  // `resource_provider`: interface to ask the browser for resources.
  //  `outbound_channel`: used to send cast messages to the receiver.
  //   `inbound_channel`: used to receiver cast messages from the receiver.
  //    `io_task_runner`: used to interact with the GPU through Viz. This arg
  //                      must be passed to enable hardware encoding.
  OpenscreenSessionHost(
      mojom::SessionParametersPtr session_params,
      const gfx::Size& max_resolution,
      mojo::PendingRemote<mojom::SessionObserver> observer,
      mojo::PendingRemote<mojom::ResourceProvider> resource_provider,
      mojo::PendingRemote<mojom::CastMessageChannel> outbound_channel,
      mojo::PendingReceiver<mojom::CastMessageChannel> inbound_channel,
      scoped_refptr<base::SingleThreadTaskRunner> io_task_runner,
      base::OnceClosure deletion_cb);

  ~OpenscreenSessionHost() override;

  // Initializes some of the asynchronous components of the session host, such
  // as access to the GPU. Must be called before negotiation of a session
  // can begin.
  using AsyncInitializedCallback = base::OnceCallback<void()>;
  void AsyncInitialize(AsyncInitializedCallback done_cb = {});

  // SenderSession::Client overrides.
  void OnNegotiated(const openscreen::cast::SenderSession* session,
                    openscreen::cast::SenderSession::ConfiguredSenders senders,
                    Recommendations capture_recommendations) override;
  void OnCapabilitiesDetermined(
      const openscreen::cast::SenderSession* session,
      openscreen::cast::RemotingCapabilities capabilities) override;
  void OnError(const openscreen::cast::SenderSession* session,
               const openscreen::Error& error) override;

  void RequestRefreshFrame();
  void CreateVideoEncodeAccelerator(
      media::cast::ReceiveVideoEncodeAcceleratorCallback callback);

  // MediaRemoter::Client overrides.
  void ConnectToRemotingSource(
      mojo::PendingRemote<media::mojom::Remoter> remoter,
      mojo::PendingReceiver<media::mojom::RemotingSource> source_receiver)
      override;
  void RequestRemotingStreaming() override;
  void RestartMirroringStreaming() override;
  std::unique_ptr<media::mojom::RemotingDataStreamSender>
  CreateRemotingDataStreamSender(
      bool is_audio,
      mojo::ScopedDataPipeConsumerHandle pipe,
      mojo::PendingReceiver<media::mojom::RemotingDataStreamSender> receiver,
      base::OnceClosure error_callback) override;

  void SwitchSourceTab();

  // Callback by media::cast::VideoSender to set a new target playout delay.
  void SetTargetPlayoutDelay(base::TimeDelta playout_delay);

  base::DictValue GetMirroringStats() const;
  void SetSenderStatsForTest(const openscreen::cast::SenderStats& test_stats);

 private:
  friend class OpenscreenSessionHostTest;
  FRIEND_TEST_ALL_PREFIXES(OpenscreenSessionHostTest, ChangeTargetPlayoutDelay);
  FRIEND_TEST_ALL_PREFIXES(OpenscreenSessionHostTest, UpdateBandwidthEstimate);

  using SupportedProfiles = media::VideoEncodeAccelerator::SupportedProfiles;

  // Called when the GPU is either set up or determined to be unavailable due
  // to software rendering being used.
  void OnAsyncInitialized(const SupportedProfiles& profiles);

  // Called when the MirroringGpuFactoriesFactory has fully initialized and
  // the GPU channel token and route ID are available.
  void OnGpuFactoriesConfigured(const base::UnguessableToken& channel_token,
                                int32_t route_id);

  // Notify `observer_` that error occurred and close the session.
  //
  // NOTE: since this method is used with base::Callback, it takes ownership
  // of the `message` to avoid lifetime issues, especially when posted to a task
  // runner.
  void ReportAndLogError(mojom::SessionError error, std::string message);

  // Stops the current streaming session. If not called from StopSession(), a
  // new streaming session will start later after exchanging OFFER/ANSWER
  // messages with the receiver. This could happen any number of times before
  // StopSession() shuts down everything permanently.
  void StopStreaming();

  // Shuts down the entire mirroring session.
  void StopSession();

  // Helper method for taking the recommendations given by the Open Screen
  // library and applying them to the given audio and video configs.
  void SetConstraints(
      const Recommendations& recommendations,
      std::optional<media::cast::FrameSenderConfig>& audio_config,
      std::optional<media::cast::FrameSenderConfig>& video_config);

  // Sends a request to create an audio input stream through the Audio Service,
  // configured with the specified audio `params`. The `shared_memory_count`
  // property indicates how many equal-lengthed segments exist in the shared
  // memory buffer. Once the stream has been created, `client` is called.
  void CreateAudioStream(
      mojo::PendingRemote<mojom::AudioStreamCreatorClient> client,
      const media::AudioParameters& params,
      uint32_t shared_memory_count);

  // Callback by Audio/VideoSender to indicate encoder status change.
  void OnAudioEncoderStatus(const media::cast::FrameSenderConfig& config,
                            media::cast::OperationalStatus status);
  void OnVideoEncoderStatus(const media::cast::FrameSenderConfig& config,
                            media::cast::OperationalStatus status);

  // Callback by MirroringGpuFactoriesFactory to indicate that the
  // GPU factory was lost (and must be replaced).
  void OnGpuFactoryContextLost(const media::cast::FrameSenderConfig& config);

  // Callback by media::cast::VideoSender to report resource utilization.
  void ProcessFeedback(const media::VideoCaptureFeedback& feedback);

  // Called by media::cast::VideoSender to help determine the video bitrate.
  uint32_t GetVideoNetworkBandwidth() const;

  // Called periodically to update the `bandwidth_estimate_`.
  void UpdateBandwidthEstimate();

  // Create and send OFFER message.
  void Negotiate();
  void NegotiateMirroring();
  void NegotiateRemoting();

  // Initialize `media_remoter_` and `rpc_dispatcher_`.
  void InitMediaRemoter(
      const openscreen::cast::RemotingCapabilities& capabilities);

  // Called 5 seconds after the `media_remoter_` is initialized for Remote
  // Playabck sessions. It terminates the streaming session if remoting is not
  // started when it's called.
  void OnRemotingStartTimeout();

  // Manage audio capture. Note the media::AudioInputDevice class does not
  // support pausing and resuming.
  void StartCapturingAudio();
  void StopCapturingAudio();

  // Manage video capture. Note that while pause and resume are supported,
  // stopping video capture is accomplished by destroying the
  // `video_capture_client_` instance.
  void StartCapturingVideo();
  void PauseCapturingVideo();

  // Returns `true` if successfully restarted video capture, otherwise it
  // may need to be started again.
  bool TryResumeCapturingVideo();

  // Called to provide Open Screen with access to this host's network proxy.
  network::mojom::SocketFactory* GetSocketFactory();

  // Called to disable the given hardware codec for the remainder of the
  // session, if it has not already been disabled.
  void MaybeDenylistHardwareCodecAndRenegotiate(media::VideoCodec codec);

  // Provided by client.
  const mojom::SessionParameters session_params_;

  // State transition:
  // kInitializing
  //     |
  //     ↓
  // kMirroring <-------> kRemoting
  //     |                   |
  //     `---> kStopped <----'
  //
  // NOTE: once a session has reached a kStopped state, it cannot be
  // reinitialized or used.
  enum class State {
    // The session is initializing, and can't be used yet.
    kInitializing,

    // A mirroring streaming session is starting or started.
    kMirroring,

    // A remoting streaming session is starting or started.
    kRemoting,

    // The session is stopped due to a user request or a fatal error.
    kStopped,
  };
  State state_ = State::kInitializing;

  // Informed of changes to session state.
  mojo::Remote<mojom::SessionObserver> observer_;

  // Provides a variety of instances, such as the current network context.
  mojo::Remote<mojom::ResourceProvider> resource_provider_;

  // Implements an Open Screen message port and wraps inbound and outbound mojom
  // channels.
  OpenscreenMessagePort message_port_;

  // Utility object for logging.
  MirroringLogger logger_;

  // Used to initialize video and audio capture clients.
  MirrorSettings mirror_settings_;

  // Used to wrap the current thread's sequenced task runner for use by Open
  // Screen.
  std::unique_ptr<openscreen_platform::TaskRunner> openscreen_task_runner_;

  // Used to wrap the `openscreen_task_runner` as well as the clock and
  // local endpoint for binding. Responsible for creating and binding a UDP
  // socket.
  std::unique_ptr<openscreen::cast::Environment> openscreen_environment_;

  // Takes care of handling OFFER/ANSWER negotiations, as well as querying
  // capabilities and creating openscreen::cast::Sender objects upon
  // negotiation.
  std::unique_ptr<openscreen::cast::SenderSession> session_;

  // Used to provide access to UDP sockets.
  mojo::Remote<network::mojom::SocketFactory> socket_factory_;
  bool set_socket_factory_proxy_ = false;

  // Stored as part of generating an OFFER.
  // NOTE: currently we only support Opus audio, but may provide a variety of
  // video codec configurations.
  std::optional<media::cast::FrameSenderConfig> last_offered_audio_config_;
  std::vector<media::cast::FrameSenderConfig> last_offered_video_configs_;
  bool offering_fallback_codecs_ = false;

  // Created after OFFER/ANSWER exchange succeeds.
  std::unique_ptr<media::cast::AudioSender> audio_sender_;
  std::unique_ptr<media::cast::VideoSender> video_sender_;

  // The time between requests for refresh frames. If zero, no refresh frames
  // will be requested.
  base::TimeDelta refresh_interval_;

  // Requests refresh frames at a constant rate while the source is paused, up
  // to a consecutive maximum.
  base::RepeatingTimer refresh_timer_;

  // Set to true when a request for a refresh frame has been made. This is
  // cleared once the next frame is received.
  bool expecting_a_refresh_frame_{false};

  // Connects to the video capture host and launches the video capture device.
  std::unique_ptr<VideoCaptureClient> video_capture_client_;

  // True if the video encoder has been initialized. This means that any further
  // calls to change encoder status to true are reinitializations, for which
  // capture should be resumed.
  bool has_video_encoder_been_initialized_ = false;

  // Manages the clock and thread proxies for the audio sender, video sender,
  // and media remoter.
  //
  // NOTE: this is lazy initialized on the first session negotiation, and then
  // destructed only on the destruction of this class.
  scoped_refptr<media::cast::CastEnvironment> cast_environment_;

  // Called when audio is successfully captured by `audio_input_device_`.
  std::unique_ptr<AudioCapturingCallback> audio_capturing_callback_;

  // Captures audio samples from the resourceprovider-created audio stream.
  scoped_refptr<media::AudioInputDevice> audio_input_device_;

  // Used as an interface for the media remoter to send RPC messages. Created
  // when a successful capabilities response arrives.
  std::unique_ptr<RpcDispatcher> rpc_dispatcher_;

  // Manages remoting content to the Cast Receiver. Created when a successful
  // capabilities response arrives.
  std::unique_ptr<MediaRemoter> media_remoter_;

  // Stored between OnNegotiated() and CreateRemotingDataStreamSender() calls
  // during a remoting session.
  struct RemotingStreamData {
    RemotingStreamData(
        std::unique_ptr<openscreen::cast::Sender> audio_sender,
        std::unique_ptr<openscreen::cast::Sender> video_sender,
        std::optional<media::cast::FrameSenderConfig> audio_config,
        std::optional<media::cast::FrameSenderConfig> video_config);
    ~RemotingStreamData();

    std::unique_ptr<openscreen::cast::Sender> audio_sender;
    std::unique_ptr<openscreen::cast::Sender> video_sender;
    std::optional<media::cast::FrameSenderConfig> audio_config;
    std::optional<media::cast::FrameSenderConfig> video_config;
  };
  std::unique_ptr<RemotingStreamData> remoting_stream_data_;

  // GPU specific properties, used to indicate whether HW encoding should be
  // used and to help initialize it if enabled.
  std::unique_ptr<viz::Gpu> gpu_;
  SupportedProfiles supported_profiles_;
  mojo::Remote<media::mojom::VideoEncodeAcceleratorProvider> vea_provider_;
  std::optional<MirroringGpuFactoriesFactory::UniquePtr> gpu_factories_factory_;
  std::vector<media::cast::ReceiveVideoEncodeAcceleratorCallback>
      pending_vea_requests_;
  base::UnguessableToken channel_token_;
  int32_t route_id_ = 0;

  // Called when the session host has fully initialized.
  AsyncInitializedCallback initialized_cb_;

  // Used to periodically update the currently used bandwidth estimate.
  base::RepeatingTimer bandwidth_update_timer_;

  // Used to override getting the bandwidth from the session. Setting to a
  // positive value causes the session's bandwidth estimation to not be called.
  int forced_bandwidth_estimate_for_testing_ = 0;

  // The portion of the bandwidth estimate that is currently available for use.
  // Note that the actual bandwidth will be effectively capped at the sum of the
  // current video and audio bitrates.
  uint32_t usable_bandwidth_ = kDefaultBitrate;

  // Indicate whether we're in the middle of switching tab sources.
  bool switching_tab_source_ = false;
  // This timer is used to stop the session in case Remoting is not started
  // before timeout. The timer is stopped when Remoting session successfully
  // starts.
  base::OneShotTimer remote_playback_start_timer_;

  // An optional stats client for fetching quality statistics from an Openscreen
  // casting session.
  std::unique_ptr<OpenscreenStatsClient> stats_client_;

  // Callback invoked once this instance and all of its resources are released.
  base::OnceClosure deletion_cb_;

  void InsertVideoFrame(scoped_refptr<media::VideoFrame> video_frame);
  void OnRefreshTimerFired();

  // Ensures that this class is accessed on a single sequence.
  SEQUENCE_CHECKER(sequence_checker_);

  // Used in callbacks executed on task runners, such as by RtpStream.
  base::WeakPtrFactory<OpenscreenSessionHost> weak_factory_{this};
};

}  // namespace mirroring

#endif  // COMPONENTS_MIRRORING_SERVICE_OPENSCREEN_SESSION_HOST_H_