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

#ifndef CHROME_BROWSER_MEDIA_AUDIO_PROCESS_ML_MODEL_FORWARDER_H_
#define CHROME_BROWSER_MEDIA_AUDIO_PROCESS_ML_MODEL_FORWARDER_H_

#include <memory>
#include <optional>

#include "base/containers/flat_map.h"
#include "base/files/file.h"
#include "base/files/file_path.h"
#include "base/memory/raw_ptr.h"
#include "base/memory/scoped_refptr.h"
#include "base/memory/weak_ptr.h"
#include "base/task/sequenced_task_runner.h"
#include "base/types/optional_ref.h"
#include "components/optimization_guide/core/delivery/optimization_guide_model_provider.h"
#include "components/optimization_guide/core/delivery/optimization_target_model_observer.h"
#include "components/optimization_guide/proto/models.pb.h"
#include "mojo/public/cpp/bindings/remote.h"
#include "services/audio/public/mojom/ml_model_manager.mojom.h"

namespace optimization_guide {
struct ModelInfo;
}  // namespace optimization_guide

class PrefService;

// Propagates ML models from the Optimization Guide to the audio process.
// Currently supports residual echo estimation and voice isolation denoising.
//
// Does nothing until both an Optimization Guide model provider has been set and
// an audio input stream has been opened. Then, subscribes to models from the
// Optimization Guide and forwards them to the audio process.
//
// NOTE: This class only forwards models to the audio process, i.e., not when
// the audio service is running as a part of the browser process. Models are
// currently not expected to be used when running in the browser process.
//
// Lives on the UI thread.
class AudioProcessMlModelForwarder {
 public:
  using WrappedFilePtr = std::unique_ptr<base::File, base::OnTaskRunnerDeleter>;

  // Default factory function. Monitors audio service process launches via
  // global APIs. `pref_service` may be null.
  static std::unique_ptr<AudioProcessMlModelForwarder> Create(
      PrefService* pref_service);

  // Testing factory function. Expects audio service process launches to be
  // signaled by calls to OnAudioProcessLaunched().
  static std::unique_ptr<AudioProcessMlModelForwarder>
  CreateWithoutAudioProcessObserverForTesting(PrefService* pref_service);

  ~AudioProcessMlModelForwarder();

  // Set the Optimization Guide model provider. May only be called once. The
  // model provider must outlive the AudioProcessMlModelForwarder.
  // Requires base::Threadpool to be initialized.
  void Initialize(
      optimization_guide::OptimizationGuideModelProvider& model_provider);

  // Needs to be called with a remote for the audio service, in order to forward
  // model updates. May be called more than once to set a new remote, e.g., on
  // service restarts.
  void OnAudioProcessLaunched(
      mojo::Remote<audio::mojom::MlModelManager> ml_model_manager);

  bool HasPendingTasksForTesting() const {
    for (const auto& [_, forwarder] : model_forwarders_) {
      if (forwarder->HasPendingTasksForTesting()) {
        return true;
      }
    }
    return false;
  }
  bool HasBoundAudioProcessRemoteForTesting() const {
    return audio_process_model_manager_.is_bound();
  }
  bool HasModelForTesting() const {
    for (const auto& [_, forwarder] : model_forwarders_) {
      if (forwarder->HasModelForTesting()) {
        return true;
      }
    }
    return false;
  }
  void FlushForTesting() {
    if (audio_process_model_manager_.is_bound()) {
      audio_process_model_manager_.FlushForTesting();
    }
  }

  // Signal that an audio capture stream has been opened. Media may not yet be
  // flowing, but permission checks have concluded successfully.
  void OnAudioCaptureStarted();

 private:
  class AudioCaptureRequestObserver;
  class AudioProcessObserver;

  // Manages the subscription, loading, and forwarding of a single ML model
  // type (specified by `OptimizationTarget`).
  //
  // SingleModelForwarder registers as an observer with the Optimization Guide
  // for a specific target. When a new model is available, it receives the
  // update on the UI thread, loads/opens the model file asynchronously on a
  // background task runner, and then forwards the opened file to the audio
  // process via the owner's mojo remote.
  //
  // Lives on the UI thread.
  class SingleModelForwarder
      : public optimization_guide::OptimizationTargetModelObserver {
   public:
    SingleModelForwarder(optimization_guide::proto::OptimizationTarget target,
                         audio::mojom::MlModelType mojo_type,
                         AudioProcessMlModelForwarder* owner);
    ~SingleModelForwarder() override;

    void Initialize(
        optimization_guide::OptimizationGuideModelProvider* model_provider,
        scoped_refptr<base::SequencedTaskRunner> background_task_runner);

    // NOTE: OnModelUpdated() may be called immediately upon registering, even
    // within the call to MaybeRegisterModelObserver().
    void MaybeRegisterModelObserver(bool audio_input_stream_creation_observed);

    void MaybeSendModelToAudioProcess();

    // Stops any ongoing loading of models.
    void CancelModelLoadingTasks();
    bool HasPendingTasksForTesting() const {
      return weak_factory_.HasWeakPtrs();
    }
    bool HasModelForTesting() const { return !model_path_.empty(); }

   private:
    void OnModelUpdated(
        optimization_guide::proto::OptimizationTarget optimization_target,
        base::optional_ref<const optimization_guide::ModelInfo> model_info)
        override;

    // Continuation for MaybeSendModelToAudioProcess(), expecting either a
    // nullptr or an open, valid model file.
    void OnModelFileOpened(WrappedFilePtr file);

    SEQUENCE_CHECKER(sequence_checker_);

    const optimization_guide::proto::OptimizationTarget target_;
    const audio::mojom::MlModelType mojo_type_;

    const raw_ptr<AudioProcessMlModelForwarder> owner_;

    // Task runner for loading model files.
    scoped_refptr<base::SequencedTaskRunner> background_task_runner_;

    // Latest update from the model provider. Empty if no path received yet or
    // if a nullopt model update has been received.
    base::FilePath model_path_;

    // Handles registration / deregistration with the Optimization Guide.
    std::optional<optimization_guide::OptimizationGuideModelProviderObservation>
        model_observation_;
    base::WeakPtrFactory<SingleModelForwarder> weak_factory_{this};
  };

  // If `audio_process_observer` is null, the forwarder does not handle audio
  // service process monitoring internally. See Create*() for details.
  AudioProcessMlModelForwarder(
      std::unique_ptr<AudioProcessObserver> audio_process_observer,
      PrefService* pref_service);

  // Register model observation if the audio process has been launched and a
  // model provider is available.
  void MaybeRegisterModelObservers();

  void MaybeSendModelsToAudioProcess();

  SEQUENCE_CHECKER(sequence_checker_);

  // Signals when the audio process is ready to start receiving model updates.
  const std::unique_ptr<AudioProcessObserver> audio_process_observer_;

  // Used for checking and updating the last observed input stream creation
  // time, in order to preload models where they are likely to be needed.
  raw_ptr<PrefService> pref_service_;

  // Handle for passing models to the audio process.
  mojo::Remote<audio::mojom::MlModelManager> audio_process_model_manager_;

  // True if and only if an audio input stream was recently created. Used to
  // gate registering model observations with the Optimization Guide.
  bool audio_input_stream_creation_observed_ = false;

  // Observes creation of audio capture streams in order to detect when models
  // are likely to be needed.
  std::unique_ptr<AudioCaptureRequestObserver> audio_capture_request_observer_;

  base::flat_map<audio::mojom::MlModelType,
                 std::unique_ptr<SingleModelForwarder>>
      model_forwarders_;
};

#endif  // CHROME_BROWSER_MEDIA_AUDIO_PROCESS_ML_MODEL_FORWARDER_H_