// Copyright 2011 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
//
// This abstract class loads a client-side model and lets you compute a phishing
// score for a set of previously extracted features.  The phishing score
// corresponds to the probability that the features are indicative of a phishing
// site.
//
// See features.h for a list of features that are currently used.

#ifndef COMPONENTS_SAFE_BROWSING_CORE_COMMON_PHISHING_CLASSIFIER_SCORER_H_
#define COMPONENTS_SAFE_BROWSING_CORE_COMMON_PHISHING_CLASSIFIER_SCORER_H_

#include <stddef.h>
#include <stdint.h>

#include <string>

#include "base/files/file.h"
#include "base/files/memory_mapped_file.h"
#include "base/functional/callback.h"
#include "base/memory/raw_ptr.h"
#include "base/memory/read_only_shared_memory_region.h"
#include "base/memory/weak_ptr.h"
#include "base/observer_list.h"
#include "base/sequence_checker.h"
#include "base/task/sequenced_task_runner.h"
#include "build/build_config.h"
#include "components/safe_browsing/core/common/fbs/client_model_generated.h"
#include "components/safe_browsing/core/common/proto/client_model.pb.h"
#include "components/safe_browsing/core/common/proto/csd.pb.h"
#include "third_party/skia/include/core/SkBitmap.h"

namespace gfx {
class Image;
}

namespace safe_browsing {

// Enum used to keep stats about the status of the Scorer creation.
// These values are persisted to logs. Entries should not be renumbered and
// numeric values should never be reused.
enum ScorerCreationStatus {
  SCORER_SUCCESS = 0,
  SCORER_FAIL_MODEL_OPEN_FAIL = 1,       // Not used anymore
  SCORER_FAIL_MODEL_FILE_EMPTY = 2,      // Not used anymore
  SCORER_FAIL_MODEL_FILE_TOO_LARGE = 3,  // Not used anymore
  SCORER_FAIL_MODEL_PARSE_ERROR = 4,
  SCORER_FAIL_MODEL_MISSING_FIELDS = 5,
  SCORER_FAIL_MAP_VISUAL_TFLITE_MODEL = 6,
  SCORER_FAIL_FLATBUFFER_INVALID_REGION = 7,
  SCORER_FAIL_FLATBUFFER_INVALID_MAPPING = 8,
  SCORER_FAIL_FLATBUFFER_FAILED_VERIFY = 9,
  SCORER_FAIL_FLATBUFFER_BAD_INDICES_OR_FIELDS = 10,  // Not used anymore
  SCORER_FAIL_FLATBUFFER_INVALID_IMAGE_EMBEDDING_TFLITE_MODEL = 11,
  SCORER_STATUS_MAX  // Always add new values before this one.
};

// Scorer methods are virtual to simplify mocking of this class,
// and to allow inheritance.
class Scorer {
 public:
  virtual ~Scorer();
  // Most clients should use the factory method.  This constructor is public
  // to allow for mock implementations.
  Scorer();

  // Factory method which creates a new Scorer object by parsing the given
  // flatbuffer or tflite model. If parsing fails this method returns NULL.
  // Use this only if region is valid.
  static std::unique_ptr<Scorer> Create(base::ReadOnlySharedMemoryRegion region,
                                        base::File visual_tflite_model);

  static std::unique_ptr<Scorer> Create(int classification_input_width,
                                        int classification_input_height,
                                        base::File tflite_visual_model);

  static std::unique_ptr<Scorer> CreateScorerWithImageEmbeddingModel(
      base::ReadOnlySharedMemoryRegion region,
      base::File visual_tflite_model,
      base::File image_embedding_model);

  static std::unique_ptr<Scorer> CreateScorerWithImageEmbeddingModel(
      int classification_input_width,
      int classification_input_height,
      base::File tflite_visual_model,
      int image_embedding_input_width,
      int image_embedding_input_height,
      base::File image_embedding_model);

  void AttachImageEmbeddingModel(base::File image_embedding_model);

  void AttachImageEmbeddingModel(int image_embedding_input_width,
                                 int image_embedding_input_height,
                                 base::File image_embedding_model);

  // This method applies the TfLite visual model to the given image for image
  // classification. It asynchronously returns the list of scores for each
  // category, in the same order as `tflite_thresholds()`.
  virtual void ApplyVisualTfLiteModel(
      const gfx::Image& image,
      base::OnceCallback<void(std::vector<double>)> callback) const;

  // This method applies the TfLite visual model to the given image for
  // image embedding. It asynchronously returns an ImageFeatureEmbedding object
  // which contains a vector of floats which is the feature vector result from
  // the Image Embedder process.
  virtual void ApplyVisualTfLiteModelImageEmbedding(
      const gfx::Image& image,
      base::OnceCallback<void(ImageFeatureEmbedding)> callback) const;

  // Returns true if a valid visual TFLite flatbuffer model is available.
  bool HasVisualTfLiteModel() const;

  // Returns the version of the visual TFLite model.
  int tflite_model_version() const;

  // Returns the image embedding model version.
  int image_embedding_tflite_model_version() const;

  // Returns the thresholds configured for the visual TFLite model categories.
  const google::protobuf::RepeatedPtrField<TfLiteModelMetadata::Threshold>&
  tflite_thresholds() const;

  // Set the dimensions for image classification.
  void SetClassificationDimensions(int classification_input_width,
                                   int classification_input_height);

  // Set the dimensions for image embedding.
  void SetImageEmbeddingDimensions(int image_embedding_input_width,
                                   int image_embedding_input_height);

  // Disable copy and move.
  Scorer(const Scorer&) = delete;
  Scorer& operator=(const Scorer&) = delete;

 private:
  friend class PhishingScorerTest;

  // Apply the tflite model to the bitmap. The scores are returned by running
  // `callback` on the provided `callback_task_runner`. This is expected to be
  // run on a helper thread.
  static void ApplyVisualTfLiteModelHelper(
      const SkBitmap& bitmap,
      int input_width,
      int input_height,
      std::string model_data,
      scoped_refptr<base::SequencedTaskRunner> callback_task_runner,
      base::OnceCallback<void(std::vector<double>)> callback);

  // Apply the TfLite model to the bitmap. The ImageFeatureEmbedding object is
  // returned by running the `callback` provided by `callback_task_runner`.
  // This is expected to be run on a helper thread.
  static void ApplyImageEmbeddingTfLiteModelHelper(
      const SkBitmap& bitmap,
      int input_width,
      int input_height,
      const std::string& model_data,
      scoped_refptr<base::SequencedTaskRunner> callback_task_runner,
      base::OnceCallback<void(ImageFeatureEmbedding)> callback);

  // Unowned. Points within flatbuffer_mapping_ and should not be free()d.
  // It remains valid till flatbuffer_mapping_ is valid and should be reassigned
  // if the mapping is updated.
  raw_ptr<const flat::ClientSideModel> flatbuffer_model_;
  base::ReadOnlySharedMemoryMapping flatbuffer_mapping_;
  google::protobuf::RepeatedPtrField<TfLiteModelMetadata::Threshold>
      thresholds_;
  base::MemoryMappedFile image_embedding_model_;
  base::MemoryMappedFile visual_tflite_model_;

  int classification_input_width_;
  int classification_input_height_;

  int image_embedding_input_width_;
  int image_embedding_input_height_;

  SEQUENCE_CHECKER(sequence_checker_);
  base::WeakPtrFactory<Scorer> weak_ptr_factory_{this};
};

#if !BUILDFLAG(IS_IOS)
// A small wrapper around a Scorer that allows callers to observe for changes in
// the model.
class ScorerStorage {
 public:
  static ScorerStorage* GetInstance();

  class Observer : public base::CheckedObserver {
   public:
    virtual void OnScorerChanged() = 0;
  };

  ScorerStorage();
  ~ScorerStorage();
  ScorerStorage(const ScorerStorage&) = delete;
  ScorerStorage& operator=(const ScorerStorage&) = delete;

  void SetScorer(std::unique_ptr<Scorer> scorer);
  Scorer* GetScorer() const;
  // We will clear the scorer in situations where the OptimizationGuide server
  // provides a null model to replace a bad model on disk.
  void ClearScorer();

  void AddObserver(Observer* observer);
  void RemoveObserver(Observer* observer);

 private:
  std::unique_ptr<Scorer> scorer_;
  base::ObserverList<Observer> observers_;
};
#endif  // !BUILDFLAG(IS_IOS)

}  // namespace safe_browsing

#endif  // COMPONENTS_SAFE_BROWSING_CORE_COMMON_PHISHING_CLASSIFIER_SCORER_H_