// Copyright 2023 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_SAFE_BROWSING_CORE_BROWSER_CLIENT_SIDE_PHISHING_MODEL_H_
#define COMPONENTS_SAFE_BROWSING_CORE_BROWSER_CLIENT_SIDE_PHISHING_MODEL_H_

#include <memory>

#include "base/callback_list.h"
#include "base/containers/span.h"
#include "base/files/file.h"
#include "base/gtest_prod_util.h"
#include "base/memory/raw_ptr.h"
#include "base/memory/read_only_shared_memory_region.h"
#include "base/memory/weak_ptr.h"
#include "base/sequence_checker.h"
#include "base/task/sequenced_task_runner.h"
#include "base/thread_annotations.h"
#include "components/optimization_guide/core/delivery/optimization_target_model_observer.h"
#include "components/safe_browsing/core/browser/csd_model_type.h"
#include "components/safe_browsing/core/common/fbs/client_model_generated.h"
#include "components/safe_browsing/core/common/proto/client_model.pb.h"
#include "third_party/tflite_support/src/tensorflow_lite_support/cc/task/vision/proto/embeddings.pb.h"

namespace optimization_guide {
class OptimizationGuideModelProvider;
}  // namespace optimization_guide

namespace safe_browsing {

// Holds an embedding we are targeting.
struct TargetEmbedding {
  TargetEmbedding(tflite::task::vision::FeatureVector embedding,
                  float threshold);
  tflite::task::vision::FeatureVector embedding;
  float threshold;
};

// This holds the currently active client side phishing detection model.
//
// The data to populate it is fetched periodically from Google to get the most
// up-to-date model. We assume it is updated at most every few hours.
//
// This class lives on UI thread and can only be called there. In particular
// GetModelStr() returns a string reference, which assumes the string won't be
// used and updated at the same time.

class ClientSidePhishingModel
    : public optimization_guide::OptimizationTargetModelObserver {
 public:
  ClientSidePhishingModel(
      optimization_guide::OptimizationGuideModelProvider* opt_guide,
      scoped_refptr<base::SequencedTaskRunner> ui_task_runner);

  ~ClientSidePhishingModel() override;

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

  // Enhanced Safe Browsing users receive the image classifier and embedding
  // model.
  void SubscribeToImageEmbedderOptimizationGuide();
  void SubscribeToImageClassifierOptimizationGuide();

  void UnsubscribeToImageEmbedderOptimizationGuide();
  void UnsubscribeToImageClassifierOptimizationGuide();

  // Register a callback to be notified whenever the model changes. All
  // notifications will occur on the UI thread.
  base::CallbackListSubscription RegisterCallback(
      base::RepeatingCallback<void()> callback);

  // Returns whether we currently have a model.
  bool IsEnabled() const;

  // Returns the hash of an embedding.
  static std::string GetHashFromEmbedding(const std::vector<float>& embedding);

  // Returns model type (flatbuffer or none).
  CSDModelType GetModelType() const;

  // Returns the shared memory region for the flatbuffer.
  base::ReadOnlySharedMemoryRegion GetModelSharedMemoryRegion() const;

  const base::File& GetVisualTfLiteModel() const;

  const base::File& GetImageEmbeddingModel() const;

  bool HasImageEmbeddingModel();

  bool IsModelMetadataImageEmbeddingVersionMatching();

  int GetTriggerModelVersion();
  int GetImageEmbeddingModelVersion();

  void SetVisualTfLiteModelForTesting(base::File file);
  // Overrides model type.
  void SetModelTypeForTesting(CSDModelType model_type);
  // Removes mapping.
  void ClearMappedRegionForTesting();
  // Get flatbuffer memory span.
  base::span<uint8_t> GetFlatBufferMemorySpanForTesting();
  // Notifies all the callbacks of a change in model.
  void NotifyCallbacksOfUpdateForTesting();
  // Sets a callback to be run when model update is complete (success or
  // failure).
  void SetModelDoneCallbackForTesting(base::OnceClosure callback);

  const std::vector<TfLiteModelMetadata::Threshold>&
  GetVisualTfLiteModelThresholds() const;

  const std::vector<TargetEmbedding>& GetTargetImageEmbeddings() const;

  void SetTargetImageEmbeddingsForTesting(
      std::vector<TargetEmbedding> target_embeddings);

  // This function is used to override internal model for testing in
  // client_side_phishing_model_unittest
  void MaybeOverrideModel();

  void OnModelAndVisualTfLiteFileLoaded(
      std::optional<optimization_guide::proto::Any> model_metadata,
      std::pair<std::string, base::File> model_and_tflite);

  void OnImageEmbeddingModelFileAndEmbeddingListLoaded(
      std::optional<optimization_guide::proto::Any> model_metadata,
      std::pair<base::File, std::optional<EmbeddingList>> model_and_list);

  void SetModelAndVisualTfLiteForTesting(
      const base::FilePath& model_file_path,
      const base::FilePath& visual_tf_lite_model_path);

  // Updates the internal model string, when one is received from testing in
  // client_side_phishing_model_unittest
  void SetModelStringForTesting(const std::string& model_str,
                                base::File visual_tflite_model);

  bool IsSubscribedToImageEmbeddingModelUpdates();
  bool IsSubscribedToImageClassifierModelUpdates();

  int GetClassificationInputWidth();
  int GetClassificationInputHeight();
  int GetImageEmbeddingInputWidth();
  int GetImageEmbeddingInputHeight();

 private:
  static const int kInitialClientModelFetchDelayMs;

  void NotifyCallbacksOnUI();

  // Callback when the file overriding the model has been read in
  // client_side_phishing_model_unittest
  void OnGetOverridenModelData(
      CSDModelType model_type,
      std::pair<std::string, base::File> model_and_tflite);

  // The list of callbacks to notify when a new model is ready. Guarded by
  // sequence_checker_. Will always be notified on the UI thread.
  base::RepeatingCallbackList<void()> callbacks_
      GUARDED_BY_CONTEXT(sequence_checker_);

  // Model protobuf string. Guarded by sequence_checker_.
  std::string model_str_ GUARDED_BY_CONTEXT(sequence_checker_);

  // Visual TFLite model file. Guarded by sequence_checker_.
  std::optional<base::File> visual_tflite_model_
      GUARDED_BY_CONTEXT(sequence_checker_);

  // Image Embedding TfLite model file. Guarded by sequence_checker_.
  std::optional<base::File> image_embedding_model_
      GUARDED_BY_CONTEXT(sequence_checker_);

  // Thresholds in visual TFLite model file to be used for comparison after
  // visual classification
  std::vector<TfLiteModelMetadata::Threshold> thresholds_;

  // Model type as inferred by feature flag. Guarded by sequence_checker_.
  CSDModelType model_type_ GUARDED_BY_CONTEXT(sequence_checker_) =
      CSDModelType::kNone;

  // MappedReadOnlyRegion where the flatbuffer has been copied to. Guarded by
  // sequence_checker_.
  base::MappedReadOnlyRegion mapped_region_
      GUARDED_BY_CONTEXT(sequence_checker_) = base::MappedReadOnlyRegion();

  FRIEND_TEST_ALL_PREFIXES(ClientSidePhishingModelTest, CanOverrideWithFlag);

  // Optimization Guide service that provides the client side detection
  // model files for this service. Optimization Guide Service is a
  // BrowserContextKeyedServiceFactory and should not be used after Shutdown
  raw_ptr<optimization_guide::OptimizationGuideModelProvider> opt_guide_;

  // These two integer values will be set from reading the metadata specified
  // under each optimization target. These two are used to match the model
  // pairings properly. If the two values match, then the image embedding model
  // will be sent to the renderer process along with the trigger models. They do
  // not reflect any versions used in the model file itself.
  std::optional<int> trigger_model_opt_guide_metadata_image_embedding_version_;
  std::optional<int>
      embedding_model_opt_guide_metadata_image_embedding_version_;

  // These values are set from a version set in the model file's metadata. This
  // value will be used to send to the CSD service class so that it can be added
  // to the debugging metadata so that we can understand what version has been
  // sent to the renderer. In addition, whenever the respective TfLite library
  // calls are completed, the model of the version used is attached to the ping.
  std::optional<int> trigger_model_version_;
  std::optional<int> image_embedding_model_version_;

  std::optional<int> classification_input_width_;
  std::optional<int> classification_input_height_;
  std::optional<int> img_embedding_input_width_;
  std::optional<int> img_embedding_input_height_;

  scoped_refptr<base::SequencedTaskRunner> ui_task_runner_;
  scoped_refptr<base::SequencedTaskRunner> background_task_runner_;

  // If the users subscribe to ESB, the code will add an observer to the
  // OptimizationGuide service for the image embedder and classifier model. We
  // can choose to remove the observer, but it will be on the list to be
  // removed, and not removed instantly. Therefore, if the user subscribes,
  // unsubscribes, and re-subscribes again in very quick succession, the code
  // will crash because the DCHECK fails, indicating that the observer is added
  // already. Therefore, this will be a one time use flag.
  bool subscribed_to_image_embedder_ = false;
  bool subscribed_to_image_classifier_ = false;

  SEQUENCE_CHECKER(sequence_checker_);

  base::TimeTicks beginning_time_;

  // The image embedding targets to evaluate pages against.
  std::vector<TargetEmbedding> target_image_embeddings_;

  base::OnceClosure model_updated_callback_for_testing_
      GUARDED_BY_CONTEXT(sequence_checker_);

  base::WeakPtrFactory<ClientSidePhishingModel> weak_ptr_factory_{this};
};

}  // namespace safe_browsing

#endif  // COMPONENTS_SAFE_BROWSING_CORE_BROWSER_CLIENT_SIDE_PHISHING_MODEL_H_