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

#include "components/optimization_guide/core/delivery/prediction_manager.h"

#include <memory>
#include <utility>
#include <vector>

#include "base/containers/contains.h"
#include "base/containers/fixed_flat_set.h"
#include "base/containers/flat_map.h"
#include "base/containers/flat_set.h"
#include "base/containers/flat_tree.h"
#include "base/functional/callback.h"
#include "base/metrics/histogram_functions.h"
#include "base/metrics/histogram_macros.h"
#include "base/metrics/histogram_macros_local.h"
#include "base/observer_list.h"
#include "base/path_service.h"
#include "base/sequence_checker.h"
#include "base/strings/string_number_conversions.h"
#include "base/task/sequenced_task_runner.h"
#include "base/task/single_thread_task_runner.h"
#include "base/task/thread_pool.h"
#include "base/time/time.h"
#include "base/uuid.h"
#include "components/optimization_guide/core/delivery/model_info.h"
#include "components/optimization_guide/core/delivery/model_util.h"
#include "components/optimization_guide/core/delivery/optimization_target_model_observer.h"
#include "components/optimization_guide/core/delivery/prediction_model_download_manager.h"
#include "components/optimization_guide/core/delivery/prediction_model_fetcher_impl.h"
#include "components/optimization_guide/core/delivery/prediction_model_override.h"
#include "components/optimization_guide/core/delivery/prediction_model_store.h"
#include "components/optimization_guide/core/optimization_guide_constants.h"
#include "components/optimization_guide/core/optimization_guide_enums.h"
#include "components/optimization_guide/core/optimization_guide_features.h"
#include "components/optimization_guide/core/optimization_guide_logger.h"
#include "components/optimization_guide/core/optimization_guide_permissions_util.h"
#include "components/optimization_guide/core/optimization_guide_prefs.h"
#include "components/optimization_guide/core/optimization_guide_switches.h"
#include "components/optimization_guide/core/optimization_guide_util.h"
#include "components/optimization_guide/optimization_guide_internals/webui/optimization_guide_internals.mojom.h"
#include "components/optimization_guide/proto/models.pb.h"
#include "components/prefs/pref_service.h"
#include "components/services/unzip/public/cpp/unzip.h"
#include "google_apis/google_api_keys.h"
#include "mojo/public/cpp/bindings/remote.h"
#include "services/network/public/cpp/shared_url_loader_factory.h"

namespace optimization_guide {

namespace {

proto::ModelCacheKey GetModelCacheKey(const std::string& locale) {
  proto::ModelCacheKey model_cache_key;
  model_cache_key.set_locale(locale);
  return model_cache_key;
}

// Util class for recording the construction and validation of a prediction
// model. The result is recorded when it goes out of scope and its destructor is
// called.
class ScopedPredictionModelConstructionAndValidationRecorder {
 public:
  explicit ScopedPredictionModelConstructionAndValidationRecorder(
      proto::OptimizationTarget optimization_target)
      : validation_start_time_(base::TimeTicks::Now()),
        optimization_target_(optimization_target) {}

  ~ScopedPredictionModelConstructionAndValidationRecorder() {
    base::UmaHistogramBoolean(
        "OptimizationGuide.IsPredictionModelValid." +
            GetStringNameForOptimizationTarget(optimization_target_),
        is_valid_);

    // Only record the timing if the model is valid and was able to be
    // constructed.
    if (is_valid_) {
      base::TimeDelta validation_latency =
          base::TimeTicks::Now() - validation_start_time_;
      base::UmaHistogramTimes(
          "OptimizationGuide.PredictionModelValidationLatency." +
              GetStringNameForOptimizationTarget(optimization_target_),
          validation_latency);
    }
  }

  void set_is_valid(bool is_valid) { is_valid_ = is_valid; }

 private:
  bool is_valid_ = true;
  const base::TimeTicks validation_start_time_;
  const proto::OptimizationTarget optimization_target_;
};

void RecordModelUpdateVersion(const proto::ModelInfo& model_info) {
  base::UmaHistogramSparse(
      "OptimizationGuide.PredictionModelUpdateVersion." +
          GetStringNameForOptimizationTarget(model_info.optimization_target()),
      model_info.version());
}

void RecordLifecycleState(proto::OptimizationTarget optimization_target,
                          ModelDeliveryEvent event) {
  base::UmaHistogramEnumeration(
      "OptimizationGuide.PredictionManager.ModelDeliveryEvents." +
          GetStringNameForOptimizationTarget(optimization_target),
      event);
}

// Returns whether models should be fetched from the
// remote Optimization Guide Service.
bool ShouldFetchModels(bool off_the_record,
                       bool component_updates_enabled,
                       bool should_check_google_api_key_configuration) {
  return features::IsRemoteFetchingEnabled() && !off_the_record &&
         features::IsModelDownloadingEnabled() && component_updates_enabled &&
         (!should_check_google_api_key_configuration ||
          google_apis::HasAPIKeyConfigured());
}

// Returns whether the model metadata proto is on the server allowlist.
bool IsModelMetadataTypeOnServerAllowlist(const proto::Any& model_metadata) {
  static const auto* const kAllowList = new base::flat_set<std::string>{
      "type.googleapis.com/"
      "google.internal.chrome.optimizationguide.v1."
      "OnDeviceTailSuggestModelMetadata",
      "type.googleapis.com/"
      "google.internal.chrome.optimizationguide.v1."
      "PageTopicsModelMetadata",
      "type.googleapis.com/"
      "google.internal.chrome.optimizationguide.v1."
      "SegmentationModelMetadata",
      "type.googleapis.com/"
      "google.privacy.webpermissionpredictions.v1."
      "WebPermissionPredictionsModelMetadata",
      "type.googleapis.com/"
      "google.internal.chrome.optimizationguide.v1."
      "ClientSidePhishingModelMetadata",
      "type.googleapis.com/"
      "lens.prime.csc.VisualSearchModelMetadata",
      "type.googleapis.com/"
      "google.internal.chrome.optimizationguide.v1."
      "OnDeviceBaseModelMetadata",
      "type.googleapis.com/"
      "google.internal.chrome.optimizationguide.v1."
      "AutofillFieldClassificationModelMetadata",
      "type.googleapis.com/"
      "google.internal.chrome.optimizationguide.v1."
      "AutocompleteScoringModelMetadata",
      "type.googleapis.com/"
      "google.privacy.webpermissionpredictions.v1."
      "WebPermissionPredictionsClientInfo",
      "type.googleapis.com/"
      "google.privacy.webpermissionpredictions.aiv3.v1."
      "PermissionsAiv3ModelMetadata"};

  return base::Contains(*kAllowList, model_metadata.type_url());
}

void RecordModelAvailableAtRegistration(
    proto::OptimizationTarget optimization_target,
    bool model_available_at_registration) {
  base::UmaHistogramBoolean(
      "OptimizationGuide.PredictionManager.ModelAvailableAtRegistration." +
          GetStringNameForOptimizationTarget(optimization_target),
      model_available_at_registration);
}

}  // namespace

PredictionManager::ModelRegistrationInfo::ModelRegistrationInfo(
    std::optional<proto::Any> metadata)
    : metadata(metadata) {}

PredictionManager::ModelRegistrationInfo::~ModelRegistrationInfo() = default;

PredictionManager::PredictionManager(
    PredictionModelStore* prediction_model_store,
    scoped_refptr<network::SharedURLLoaderFactory> url_loader_factory,
    PrefService* pref_service,
    bool off_the_record,
    const std::string& application_locale,
    OptimizationGuideLogger* optimization_guide_logger,
    ComponentUpdatesEnabledProvider component_updates_enabled_provider,
    unzip::UnzipperFactory unzipper_factory)
    : prediction_model_download_manager_(nullptr),
      prediction_model_store_(prediction_model_store),
      url_loader_factory_(url_loader_factory),
      optimization_guide_logger_(optimization_guide_logger),
      component_updates_enabled_provider_(component_updates_enabled_provider),
      unzipper_factory_(std::move(unzipper_factory)),
      prediction_model_fetch_timer_(
          pref_service,
          base::BindRepeating(
              &PredictionManager::FetchModels,
              // Its safe to use `base::Unretained(this)` here since
              // `prediction_model_fetch_timer_` is owned by `this`.
              base::Unretained(this))),
      off_the_record_(off_the_record),
      application_locale_(application_locale),
      model_cache_key_(GetModelCacheKey(application_locale_)),
      should_check_google_api_key_configuration_(
          !switches::ShouldSkipGoogleApiKeyConfigurationCheck()) {
  DCHECK(prediction_model_store_);
  LoadPredictionModels(GetRegisteredOptimizationTargets());
  LOCAL_HISTOGRAM_BOOLEAN(
      "OptimizationGuide.PredictionManager.StoreInitialized", true);
}

PredictionManager::~PredictionManager() {
  if (prediction_model_download_manager_) {
    prediction_model_download_manager_->RemoveObserver(this);
  }
}

void PredictionManager::AddObserverForOptimizationTargetModel(
    proto::OptimizationTarget optimization_target,
    const std::optional<proto::Any>& model_metadata,
    OptimizationTargetModelObserver* observer) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);

  // A limited number of targets support multiple registrations. In general
  // multiple registrations are disallowed to mitigate the risk of subtle,
  // conflicting behavior between two different uses of the same model file. If
  // adding a target to this set, please document below why it's necessary.
  constexpr auto kAllowedMultipleRegistrations =
      base::MakeFixedFlatSet<proto::OptimizationTarget>({
          // In addition to use by Translate's language detection features, this
          // model is also needed by the On-Device Model service process, and
          // ModelExecutionManager monitors for updates on its behalf.
          proto::OptimizationTarget::OPTIMIZATION_TARGET_LANGUAGE_DETECTION,
      });

  DCHECK(base::Contains(kAllowedMultipleRegistrations, optimization_target) ||
         !base::Contains(model_registration_info_map_, optimization_target));
  DCHECK(!model_metadata ||
         IsModelMetadataTypeOnServerAllowlist(*model_metadata));

  // As DCHECKS don't run in the wild, just do not register the observer if
  // something is already registered for the type. Otherwise, file reads may
  // blow up.
  if (!base::Contains(kAllowedMultipleRegistrations, optimization_target) &&
      base::Contains(model_registration_info_map_, optimization_target)) {
    DLOG(ERROR) << "Did not add observer for optimization target "
                << static_cast<int>(optimization_target)
                << " since an observer for the target was already registered ";
    return;
  }

  auto [it, registered] = model_registration_info_map_.emplace(
      std::piecewise_construct, std::forward_as_tuple(optimization_target),
      std::forward_as_tuple(model_metadata));
  DCHECK(registered ||
         base::Contains(kAllowedMultipleRegistrations, optimization_target));
  it->second.model_observers.AddObserver(observer);
  if (optimization_guide_logger_->ShouldEnableDebugLogs()) {
    OPTIMIZATION_GUIDE_LOGGER(
        optimization_guide_common::mojom::LogSource::MODEL_MANAGEMENT,
        optimization_guide_logger_)
        << "Observer added for OptimizationTarget: " << optimization_target;
  }

  // Notify observer of existing model file path.
  auto model_it = optimization_target_model_info_map_.find(optimization_target);
  if (model_it != optimization_target_model_info_map_.end()) {
    observer->OnModelUpdated(optimization_target, *model_it->second);
    if (optimization_guide_logger_->ShouldEnableDebugLogs()) {
      OPTIMIZATION_GUIDE_LOGGER(
          optimization_guide_common::mojom::LogSource::MODEL_MANAGEMENT,
          optimization_guide_logger_)
          << "OnModelFileUpdated for OptimizationTarget: "
          << optimization_target << "\nFile path: "
          << model_it->second->GetModelFilePath().AsUTF8Unsafe()
          << "\nHas metadata: " << (model_metadata ? "True" : "False");
    }
    RecordLifecycleState(optimization_target,
                         ModelDeliveryEvent::kModelDeliveredAtRegistration);
  }
  base::UmaHistogramMediumTimes(
      "OptimizationGuide.PredictionManager.RegistrationTimeSinceServiceInit." +
          GetStringNameForOptimizationTarget(optimization_target),
      !init_time_.is_null() ? base::TimeTicks::Now() - init_time_
                            : base::TimeDelta());

  if (optimization_guide_logger_->ShouldEnableDebugLogs()) {
    OPTIMIZATION_GUIDE_LOGGER(
        optimization_guide_common::mojom::LogSource::MODEL_MANAGEMENT,
        optimization_guide_logger_)
        << "Registered new OptimizationTarget: " << optimization_target;
  }

  if (ShouldFetchModels(off_the_record_,
                        component_updates_enabled_provider_.Run(),
                        should_check_google_api_key_configuration_)) {
    prediction_model_fetch_timer_.ScheduleFetchOnModelRegistration();
  }

  // Otherwise, load prediction models for any newly registered targets.
  LoadPredictionModels({optimization_target});
}

void PredictionManager::RemoveObserverForOptimizationTargetModel(
    proto::OptimizationTarget optimization_target,
    OptimizationTargetModelObserver* observer) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  auto registration_info =
      model_registration_info_map_.find(optimization_target);
  CHECK(registration_info != model_registration_info_map_.end());

  auto& observers = registration_info->second.model_observers;
  DCHECK(observers.HasObserver(observer));
  observers.RemoveObserver(observer);
  if (observers.empty()) {
    model_registration_info_map_.erase(registration_info);
  }
}

base::flat_set<proto::OptimizationTarget>
PredictionManager::GetRegisteredOptimizationTargets() const {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  base::flat_set<proto::OptimizationTarget> optimization_targets;
  for (const auto& registration_info : model_registration_info_map_) {
    optimization_targets.insert(registration_info.first);
  }
  return optimization_targets;
}

void PredictionManager::SetPredictionModelFetcherForTesting(
    std::unique_ptr<PredictionModelFetcher> prediction_model_fetcher) {
  prediction_model_fetcher_ = std::move(prediction_model_fetcher);
}

void PredictionManager::SetPredictionModelDownloadManagerForTesting(
    std::unique_ptr<PredictionModelDownloadManager>
        prediction_model_download_manager) {
  prediction_model_download_manager_ =
      std::move(prediction_model_download_manager);
}

void PredictionManager::FetchModels() {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);

  // The histogram that gets recorded here is used for integration tests that
  // pass in a model override. For simplicity, we place the recording of this
  // histogram here rather than somewhere else earlier in the session
  // initialization flow since the model engine version needs to continuously be
  // updated for the fetch.
  proto::ModelInfo base_model_info;
  // There should only be one supported model engine version at a time.
  base_model_info.add_supported_model_engine_versions(
      proto::MODEL_ENGINE_VERSION_TFLITE_2_20_1);
  // This histogram is used for integration tests. Do not remove.
  // Update this to be 10000 if/when we exceed 100 model engine versions.
  LOCAL_HISTOGRAM_COUNTS_100(
      "OptimizationGuide.PredictionManager.SupportedModelEngineVersion",
      static_cast<int>(
          *base_model_info.supported_model_engine_versions().begin()));

  if (!ShouldFetchModels(off_the_record_,
                         component_updates_enabled_provider_.Run(),
                         should_check_google_api_key_configuration_)) {
    return;
  }

  if (prediction_model_fetch_timer_.IsFirstModelFetch()) {
    DCHECK(!init_time_.is_null());
    base::UmaHistogramMediumTimes(
        "OptimizationGuide.PredictionManager.FirstModelFetchSinceServiceInit",
        base::TimeTicks::Now() - init_time_);
  }

  // Models should not be fetched if there are no optimization targets
  // registered.
  if (model_registration_info_map_.empty()) {
    return;
  }

  // We should have already created a prediction model download manager if we
  // initiated the fetching of models.
  DCHECK(prediction_model_download_manager_);
  if (prediction_model_download_manager_) {
    bool download_service_available =
        prediction_model_download_manager_->IsAvailableForDownloads();
    base::UmaHistogramBoolean(
        "OptimizationGuide.PredictionManager."
        "DownloadServiceAvailabilityBlockedFetch",
        !download_service_available);
    if (!download_service_available) {
      for (const auto& registration_info : model_registration_info_map_) {
        RecordLifecycleState(registration_info.first,
                             ModelDeliveryEvent::kDownloadServiceUnavailable);
      }
      // We cannot download any models from the server, so don't refresh them.
      return;
    }

    prediction_model_download_manager_->CancelAllPendingDownloads();
  }

  std::vector<proto::ModelInfo> models_info = std::vector<proto::ModelInfo>();
  models_info.reserve(model_registration_info_map_.size());

  // For now, we will fetch for all registered optimization targets.
  auto overrides = PredictionModelOverrides::ParseFromCommandLine(
      base::CommandLine::ForCurrentProcess());
  for (const auto& registration_info : model_registration_info_map_) {
    if (overrides.Get(registration_info.first)) {
      // Do not download models that were overridden.
      continue;
    }

    proto::ModelInfo model_info(base_model_info);
    model_info.set_optimization_target(registration_info.first);
    if (registration_info.second.metadata) {
      *model_info.mutable_model_metadata() = *registration_info.second.metadata;
    }

    auto model_it =
        optimization_target_model_info_map_.find(registration_info.first);
    if (model_it != optimization_target_model_info_map_.end()) {
      model_info.set_version(model_it->second.get()->GetVersion());
    }

    models_info.push_back(model_info);
    if (optimization_guide_logger_->ShouldEnableDebugLogs()) {
      OPTIMIZATION_GUIDE_LOGGER(
          optimization_guide_common::mojom::LogSource::MODEL_MANAGEMENT,
          optimization_guide_logger_)
          << "Fetching models for Optimization Target "
          << model_info.optimization_target();
    }
    RecordLifecycleState(registration_info.first,
                         ModelDeliveryEvent::kGetModelsRequest);
  }
  if (models_info.empty()) {
    return;
  }

  // NOTE: ALL PRECONDITIONS FOR THIS FUNCTION MUST BE CHECKED ABOVE THIS LINE.
  // It is assumed that if we proceed past here, that a fetch will at least be
  // attempted.

  if (!prediction_model_fetcher_) {
    prediction_model_fetcher_ = std::make_unique<PredictionModelFetcherImpl>(
        url_loader_factory_,
        features::GetOptimizationGuideServiceGetModelsURL());
  }

  bool fetch_initiated =
      prediction_model_fetcher_->FetchOptimizationGuideServiceModels(
          models_info, proto::CONTEXT_BATCH_UPDATE_MODELS, application_locale_,
          base::BindOnce(&PredictionManager::OnModelsFetched,
                         ui_weak_ptr_factory_.GetWeakPtr(), models_info));

  if (fetch_initiated) {
    prediction_model_fetch_timer_.NotifyModelFetchAttempt();
  }
  // Schedule the next fetch regardless since we may not have initiated a fetch
  // due to a network condition and trying in the next minute to see if that is
  // unblocked is only a timer firing and not an actual query to the server.
  prediction_model_fetch_timer_.SchedulePeriodicModelsFetch();
}

void PredictionManager::OnModelsFetched(
    const std::vector<proto::ModelInfo> models_request_info,
    std::optional<std::unique_ptr<proto::GetModelsResponse>>
        get_models_response_data) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  if (!get_models_response_data) {
    for (const auto& model_info : models_request_info) {
      RecordLifecycleState(model_info.optimization_target(),
                           ModelDeliveryEvent::kGetModelsResponseFailure);
    }
    return;
  }

  if ((*get_models_response_data)->models_size() > 0 ||
      models_request_info.size() > 0) {
    UpdatePredictionModels(models_request_info,
                           (*get_models_response_data)->models());
  }

  prediction_model_fetch_timer_.NotifyModelFetchSuccess();
  prediction_model_fetch_timer_.Stop();
  prediction_model_fetch_timer_.SchedulePeriodicModelsFetch();
}

void PredictionManager::UpdateModelMetadata(
    const proto::PredictionModel& model) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);

  // Model update is needed when download URL is set, which indicates the model
  // has changed.
  if (model.model().download_url().empty()) {
    return;
  }
  if (!model.model_info().has_model_cache_key()) {
    return;
  }
  if (optimization_guide_logger_->ShouldEnableDebugLogs()) {
    OPTIMIZATION_GUIDE_LOGGER(
        optimization_guide_common::mojom::LogSource::MODEL_MANAGEMENT,
        optimization_guide_logger_)
        << "Optimization Target: " << model.model_info().optimization_target()
        << " for locale: " << application_locale_
        << " sharing models with locale: "
        << model.model_info().model_cache_key().locale();
  }
  prediction_model_store_->UpdateModelCacheKeyMapping(
      model.model_info().optimization_target(), model_cache_key_,
      model.model_info().model_cache_key());
}

bool PredictionManager::ShouldDownloadNewModel(
    const proto::PredictionModel& model) const {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);

  // No download needed if URL is not set.
  if (model.model().download_url().empty()) {
    return false;
  }
  // Though the server set the download URL indicating the model is old or does
  // not exist in client, the same version model could exist in the store, if
  // the model is shared across different profile characteristics, based on
  // ModelCacheKey. So, download only when the same version is not found in the
  // store.
  return !prediction_model_store_->HasModelWithVersion(
      model.model_info().optimization_target(), model_cache_key_,
      model.model_info().version());
}

void PredictionManager::StartModelDownload(
    proto::OptimizationTarget optimization_target,
    const GURL& download_url) {
  // We should only be downloading models and updating the store for
  // on-the-record profiles and after the store has been initialized.
  DCHECK(prediction_model_download_manager_);
  if (!prediction_model_download_manager_) {
    return;
  }
  if (download_url.is_valid()) {
    prediction_model_download_manager_->StartDownload(download_url,
                                                      optimization_target);
    if (optimization_guide_logger_->ShouldEnableDebugLogs()) {
      OPTIMIZATION_GUIDE_LOGGER(
          optimization_guide_common::mojom::LogSource::MODEL_MANAGEMENT,
          optimization_guide_logger_)
          << "Model download started for Optimization Target: "
          << optimization_target << " download URL: " << download_url;
    }
  }
  RecordLifecycleState(optimization_target,
                       download_url.is_valid()
                           ? ModelDeliveryEvent::kDownloadServiceRequest
                           : ModelDeliveryEvent::kDownloadURLInvalid);
  base::UmaHistogramBoolean(
      "OptimizationGuide.PredictionManager.IsDownloadUrlValid." +
          GetStringNameForOptimizationTarget(optimization_target),
      download_url.is_valid());
}

void PredictionManager::MaybeDownloadOrUpdatePredictionModel(
    proto::OptimizationTarget optimization_target,
    const proto::PredictionModel& get_models_response_model,
    std::unique_ptr<proto::PredictionModel> loaded_model) {
  if (!loaded_model) {
    // Model load failed, redownload the model.
    RecordLifecycleState(optimization_target,
                         ModelDeliveryEvent::kModelLoadFailed);
    DCHECK(!get_models_response_model.model().download_url().empty());
    StartModelDownload(optimization_target,
                       GURL(get_models_response_model.model().download_url()));
    return;
  }
  prediction_model_store_->UpdateMetadataForExistingModel(
      optimization_target, model_cache_key_,
      get_models_response_model.model_info());
  OnLoadPredictionModel(optimization_target,
                        /*record_availability_metrics=*/false,
                        std::move(loaded_model));
}

void PredictionManager::UpdatePredictionModels(
    const std::vector<proto::ModelInfo>& models_request_info,
    const google::protobuf::RepeatedPtrField<proto::PredictionModel>&
        prediction_models) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);

  std::set<proto::OptimizationTarget> received_optimization_targets;
  for (const auto& model : prediction_models) {
    auto optimization_target = model.model_info().optimization_target();
    received_optimization_targets.emplace(optimization_target);
    if (!model.has_model()) {
      // We already have this updated model, so don't update in store.
      continue;
    }
    DCHECK(!model.model().download_url().empty());
    UpdateModelMetadata(model);
    if (ShouldDownloadNewModel(model)) {
      StartModelDownload(optimization_target,
                         GURL(model.model().download_url()));
      // Skip over models that have a download URL since they will be updated
      // once the download has completed successfully.
      continue;
    }

    RecordModelUpdateVersion(model.model_info());
    DCHECK(prediction_model_store_->HasModel(optimization_target,
                                             model_cache_key_));
    // Load the model from the store to see whether it is valid or not.
    prediction_model_store_->LoadModel(
        optimization_target, model_cache_key_,
        base::BindOnce(&PredictionManager::MaybeDownloadOrUpdatePredictionModel,
                       ui_weak_ptr_factory_.GetWeakPtr(), optimization_target,
                       model));
    if (optimization_guide_logger_->ShouldEnableDebugLogs()) {
      OPTIMIZATION_GUIDE_LOGGER(
          optimization_guide_common::mojom::LogSource::MODEL_MANAGEMENT,
          optimization_guide_logger_)
          << "Model Download Not Required for target: " << optimization_target
          << "\nNew Version: "
          << base::NumberToString(model.model_info().version());
    }
  }

  for (const auto& model_info : models_request_info) {
    if (received_optimization_targets.find(model_info.optimization_target()) ==
        received_optimization_targets.end()) {
      RemoveModelFromStore(
          model_info.optimization_target(),
          PredictionModelStoreModelRemovalReason::kNoModelInGetModelsResponse);
    }
  }
}

void PredictionManager::OnModelReady(const base::FilePath& base_model_dir,
                                     const proto::PredictionModel& model) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  DCHECK(model.model_info().has_version() &&
         model.model_info().has_optimization_target());

  auto overrides = PredictionModelOverrides::ParseFromCommandLine(
      base::CommandLine::ForCurrentProcess());
  if (overrides.Get(model.model_info().optimization_target())) {
    // Skip updating the model if override is present.
    return;
  }

  RecordModelUpdateVersion(model.model_info());
  RecordLifecycleState(model.model_info().optimization_target(),
                       ModelDeliveryEvent::kModelDownloaded);
  if (optimization_guide_logger_->ShouldEnableDebugLogs()) {
    OPTIMIZATION_GUIDE_LOGGER(
        optimization_guide_common::mojom::LogSource::MODEL_MANAGEMENT,
        optimization_guide_logger_)
        << "Model Files Downloaded target: "
        << model.model_info().optimization_target()
        << "\nNew Version: " +
               base::NumberToString(model.model_info().version());
  }

  // Store the received model in the store.
  prediction_model_store_->UpdateModel(
      model.model_info().optimization_target(), model_cache_key_,
      model.model_info(), base_model_dir,
      base::BindOnce(&PredictionManager::OnPredictionModelsStored,
                     ui_weak_ptr_factory_.GetWeakPtr()));

  if (model_registration_info_map_.contains(
          model.model_info().optimization_target())) {
    OnLoadPredictionModel(model.model_info().optimization_target(),
                          /*record_availability_metrics=*/false,
                          std::make_unique<proto::PredictionModel>(model));
  }
}

void PredictionManager::OnModelDownloadStarted(
    proto::OptimizationTarget optimization_target) {
  RecordLifecycleState(optimization_target,
                       ModelDeliveryEvent::kModelDownloadStarted);
}

void PredictionManager::OnModelDownloadFailed(
    proto::OptimizationTarget optimization_target) {
  RecordLifecycleState(optimization_target,
                       ModelDeliveryEvent::kModelDownloadFailure);
}

std::vector<optimization_guide_internals::mojom::DownloadedModelInfoPtr>
PredictionManager::GetDownloadedModelsInfoForWebUI() const {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  std::vector<optimization_guide_internals::mojom::DownloadedModelInfoPtr>
      downloaded_models_info;
  downloaded_models_info.reserve(optimization_target_model_info_map_.size());
  for (const auto& it : optimization_target_model_info_map_) {
    const std::string& optimization_target_name =
        optimization_guide::proto::OptimizationTarget_Name(it.first);
    const optimization_guide::ModelInfo* const model_info = it.second.get();
    auto downloaded_model_info_ptr =
        optimization_guide_internals::mojom::DownloadedModelInfo::New(
            optimization_target_name, model_info->GetVersion(),
            model_info->GetModelFilePath().AsUTF8Unsafe());
    downloaded_models_info.push_back(std::move(downloaded_model_info_ptr));
  }
  return downloaded_models_info;
}

base::flat_map<std::string, bool>
PredictionManager::GetOnDeviceSupplementaryModelsInfoForWebUI() const {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  std::vector<proto::OptimizationTarget> supp_targets = {
      proto::OptimizationTarget::OPTIMIZATION_TARGET_TEXT_SAFETY,
      proto::OptimizationTarget::OPTIMIZATION_TARGET_LANGUAGE_DETECTION};
  base::flat_map<std::string, bool> supp_models_info;
  for (const auto target : supp_targets) {
    supp_models_info[optimization_guide::proto::OptimizationTarget_Name(
        target)] = optimization_target_model_info_map_.contains(target);
  }

  return supp_models_info;
}

void PredictionManager::NotifyObserversOfNewModel(
    proto::OptimizationTarget optimization_target,
    base::optional_ref<const ModelInfo> model_info) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  auto registration_info_it =
      model_registration_info_map_.find(optimization_target);
  if (registration_info_it == model_registration_info_map_.end()) {
    return;
  }
  RecordLifecycleState(optimization_target,
                       ModelDeliveryEvent::kModelDelivered);
  for (auto& observer : registration_info_it->second.model_observers) {
    observer.OnModelUpdated(optimization_target, model_info);
  }
  if (optimization_guide_logger_->ShouldEnableDebugLogs()) {
    if (model_info.has_value()) {
      OPTIMIZATION_GUIDE_LOGGER(
          optimization_guide_common::mojom::LogSource::MODEL_MANAGEMENT,
          optimization_guide_logger_)
          << "OnModelFileUpdated for target: " << optimization_target
          << "\nFile path: " << model_info->GetModelFilePath().AsUTF8Unsafe()
          << "\nHas metadata: "
          << (model_info->GetModelMetadata() ? "True" : "False");
    } else {
      OPTIMIZATION_GUIDE_LOGGER(
          optimization_guide_common::mojom::LogSource::MODEL_MANAGEMENT,
          optimization_guide_logger_)
          << "OnModelFileUpdated for target: " << optimization_target
          << " for model removed";
    }
  }
}

void PredictionManager::OnPredictionModelsStored() {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  LOCAL_HISTOGRAM_BOOLEAN(
      "OptimizationGuide.PredictionManager.PredictionModelsStored", true);
}

void PredictionManager::MaybeInitializeModelDownloads(
    download::BackgroundDownloadService* background_download_service) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  init_time_ = base::TimeTicks::Now();

  // Create the download manager here if we are allowed to.
  if (features::IsModelDownloadingEnabled() && !off_the_record_ &&
      !prediction_model_download_manager_) {
    prediction_model_download_manager_ =
        std::make_unique<PredictionModelDownloadManager>(
            background_download_service,
            base::BindRepeating(
                &PredictionManager::GetBaseModelDirForDownload,
                // base::Unretained is safe here because the
                // PredictionModelDownloadManager is owned by `this`
                base::Unretained(this)),
            unzipper_factory_,
            base::ThreadPool::CreateSequencedTaskRunner(
                {base::MayBlock(), base::TaskPriority::BEST_EFFORT,
                 base::TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN}));
    prediction_model_download_manager_->AddObserver(this);
  }

  // Only load models if there are optimization targets registered.
  if (!model_registration_info_map_.empty() &&
      ShouldFetchModels(off_the_record_,
                        component_updates_enabled_provider_.Run(),
                        should_check_google_api_key_configuration_)) {
    prediction_model_fetch_timer_.MaybeScheduleFirstModelFetch();
  }
}

void PredictionManager::OnPredictionModelOverrideLoaded(
    proto::OptimizationTarget optimization_target,
    std::unique_ptr<proto::PredictionModel> prediction_model) {
  const bool is_available = prediction_model != nullptr;
  VLOG(0) << "Loading override for "
          << proto::OptimizationTarget_Name(optimization_target)
          << (is_available ? " succeeded" : " failed");
  OnLoadPredictionModel(optimization_target,
                        /*record_availability_metrics=*/false,
                        std::move(prediction_model));
  RecordModelAvailableAtRegistration(optimization_target, is_available);
}

void PredictionManager::LoadPredictionModels(
    const base::flat_set<proto::OptimizationTarget>& optimization_targets) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);

  auto overrides = PredictionModelOverrides::ParseFromCommandLine(
      base::CommandLine::ForCurrentProcess());
  for (proto::OptimizationTarget optimization_target : optimization_targets) {
    // Give preference to any overrides given on the command line.
    if (auto* entry = overrides.Get(optimization_target); entry) {
      base::FilePath base_model_dir =
          GetBaseModelDirForDownload(optimization_target);
      entry->BuildModel(
          base_model_dir, unzipper_factory_,
          base::BindOnce(&PredictionManager::OnPredictionModelOverrideLoaded,
                         ui_weak_ptr_factory_.GetWeakPtr(),
                         optimization_target));
      continue;
    }

    if (!prediction_model_store_->HasModel(optimization_target,
                                           model_cache_key_)) {
      RecordModelAvailableAtRegistration(optimization_target, false);
      continue;
    }
    prediction_model_store_->LoadModel(
        optimization_target, model_cache_key_,
        base::BindOnce(&PredictionManager::OnLoadPredictionModel,
                       ui_weak_ptr_factory_.GetWeakPtr(), optimization_target,
                       /*record_availability_metrics=*/true));
  }
}

void PredictionManager::OnLoadPredictionModel(
    proto::OptimizationTarget optimization_target,
    bool record_availability_metrics,
    std::unique_ptr<proto::PredictionModel> model) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  if (!model) {
    if (record_availability_metrics) {
      RecordModelAvailableAtRegistration(optimization_target, false);
    }
    return;
  }
  bool success = ProcessAndStoreLoadedModel(*model);
  DCHECK_EQ(optimization_target, model->model_info().optimization_target());
  if (record_availability_metrics) {
    RecordModelAvailableAtRegistration(optimization_target, success);
  }
  OnProcessLoadedModel(*model, success);
}

void PredictionManager::OnProcessLoadedModel(
    const proto::PredictionModel& model,
    bool success) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);

  if (success) {
    base::UmaHistogramSparse("OptimizationGuide.PredictionModelLoadedVersion." +
                                 GetStringNameForOptimizationTarget(
                                     model.model_info().optimization_target()),
                             model.model_info().version());
    return;
  }
  RemoveModelFromStore(
      model.model_info().optimization_target(),
      PredictionModelStoreModelRemovalReason::kModelLoadFailed);
}

void PredictionManager::RemoveModelFromStore(
    proto::OptimizationTarget optimization_target,
    PredictionModelStoreModelRemovalReason model_removal_reason) {
  if (prediction_model_store_->HasModel(optimization_target,
                                        model_cache_key_)) {
    prediction_model_store_->RemoveModel(optimization_target, model_cache_key_,
                                         model_removal_reason);
    NotifyObserversOfNewModel(optimization_target, std::nullopt);
  }
}

bool PredictionManager::ProcessAndStoreLoadedModel(
    const proto::PredictionModel& model) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  if (!model.model_info().has_optimization_target()) {
    return false;
  }
  if (!model.model_info().has_version()) {
    return false;
  }
  if (!model.has_model()) {
    return false;
  }
  if (!model_registration_info_map_.contains(
          model.model_info().optimization_target())) {
    return false;
  }

  ScopedPredictionModelConstructionAndValidationRecorder
      prediction_model_recorder(model.model_info().optimization_target());
  std::unique_ptr<ModelInfo> model_info = ModelInfo::Create(model);
  if (!model_info) {
    prediction_model_recorder.set_is_valid(false);
    return false;
  }

  proto::OptimizationTarget optimization_target =
      model.model_info().optimization_target();

  // See if we should update the loaded model.
  if (!ShouldUpdateStoredModelForTarget(optimization_target,
                                        model.model_info().version())) {
    return true;
  }

  // Update prediction model file if that is what we have loaded.
  if (model_info) {
    StoreLoadedModelInfo(optimization_target, std::move(model_info));
  }

  return true;
}

bool PredictionManager::ShouldUpdateStoredModelForTarget(
    proto::OptimizationTarget optimization_target,
    int64_t new_version) const {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);

  auto model_meta_it =
      optimization_target_model_info_map_.find(optimization_target);
  if (model_meta_it != optimization_target_model_info_map_.end()) {
    return model_meta_it->second->GetVersion() != new_version;
  }

  return true;
}

void PredictionManager::StoreLoadedModelInfo(
    proto::OptimizationTarget optimization_target,
    std::unique_ptr<ModelInfo> model_info) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  DCHECK(model_info);

  // Notify observers of new model file path.
  base::SingleThreadTaskRunner::GetCurrentDefault()->PostTask(
      FROM_HERE, base::BindOnce(&PredictionManager::NotifyObserversOfNewModel,
                                ui_weak_ptr_factory_.GetWeakPtr(),
                                optimization_target, *model_info));

  optimization_target_model_info_map_.insert_or_assign(optimization_target,
                                                       std::move(model_info));
}

base::FilePath PredictionManager::GetBaseModelDirForDownload(
    proto::OptimizationTarget optimization_target) {
  return prediction_model_store_->GetBaseModelDirForModelCacheKey(
      optimization_target, model_cache_key_);
}

void PredictionManager::OverrideTargetModelForTesting(
    proto::OptimizationTarget optimization_target,
    std::unique_ptr<ModelInfo> model_info) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);

  std::optional<ModelInfo> model_info_copy;
  if (model_info) {
    model_info_copy = *model_info;
  }
  optimization_target_model_info_map_.insert_or_assign(optimization_target,
                                                       std::move(model_info));

  NotifyObserversOfNewModel(optimization_target, model_info_copy);
}

}  // namespace optimization_guide