// Copyright 2024 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/feature_registry/feature_registration.h"

#include "base/feature_list.h"
#include "components/optimization_guide/core/feature_registry/enterprise_policy_registry.h"
#include "components/optimization_guide/core/feature_registry/mqls_feature_registry.h"
#include "components/optimization_guide/core/feature_registry/settings_ui_registry.h"
#include "components/optimization_guide/core/model_execution/feature_keys.h"
#include "components/optimization_guide/proto/features/common_quality_data.pb.h"
#include "components/optimization_guide/proto/model_quality_service.pb.h"
#include "components/prefs/pref_registry_simple.h"

namespace optimization_guide {

namespace prefs {

const char kComposeEnterprisePolicyAllowed[] =
    "optimization_guide.model_execution.compose_enterprise_policy_allowed";

const char kWallpaperSearchEnterprisePolicyAllowed[] =
    "optimization_guide.model_execution.wallpaper_search_enterprise_policy_"
    "allowed";

const char kHistorySearchEnterprisePolicyAllowed[] =
    "optimization_guide.model_execution.history_search_"
    "enterprise_policy_allowed";

const char kProductSpecificationsEnterprisePolicyAllowed[] =
    "optimization_guide.model_execution.tab_compare_settings_enterprise_policy";

const char kAutofillPredictionImprovementsEnterprisePolicyAllowed[] =
    "optimization_guide.model_execution.autofill_prediction_improvements_"
    "enterprise_policy_allowed";

const char kAutomatedPasswordChangeEnterprisePolicyAllowed[] =
    "optimization_guide.model_execution.password_change_submission_"
    "enterprise_policy_allowed";

const char kNotificationContentDetectionEnterprisePolicyAllowed[] =
    "optimization_guide.model_execution.notification_content_detection_"
    "enterprise_policy_allowed";

const char kBlingPrototypingEnterprisePolicyAllowed[] =
    "optimization_guide.model_execution.bling_prototyping_enterprise_policy_"
    "allowed";

const char kContextualTasksContextEnterprisePolicyAllowed[] =
    "optimization_guide.model_execution.contextual_tasks_context_enterprise_"
    "policy_allowed";

const char kGeminiAntiscamProtectionEnterprisePolicyAllowed[] =
    "optimization_guide.model_execution.gemini_antiscam_protection_enterprise_"
    "policy_allowed";

const char kFindsEnterprisePolicyAllowed[] =
    "optimization_guide.model_execution.finds_enterprise_policy_allowed";

const char kChromeSuggestionsSettings[] =
    "contextual_cueing.chrome_suggestions_settings";

const char kGeminiSettings[] = "browser.gemini_settings";

// LINT.IfChange(FindAndFillWithGeminiSettings)
const char kFindAndFillWithGeminiSettings[] =
    "autofill.personal_context.find_and_fill_with_gemini_settings";
// LINT.ThenChange(//chrome/browser/ui/android/autofill/internal/java/src/org/chromium/chrome/browser/ui/autofill/AtMemoryBottomSheetMediator.java:FindAndFillWithGeminiSettings)
}  // namespace prefs

namespace features {
BASE_FEATURE(kActorLoginMqlsLogging, base::FEATURE_ENABLED_BY_DEFAULT);

BASE_FEATURE(kComposeMqlsLogging, base::FEATURE_ENABLED_BY_DEFAULT);

BASE_FEATURE(kWallpaperSearchMqlsLogging, base::FEATURE_ENABLED_BY_DEFAULT);

BASE_FEATURE(kHistorySearchMqlsLogging, base::FEATURE_ENABLED_BY_DEFAULT);

BASE_FEATURE(kProductSpecificationsMqlsLogging,
             base::FEATURE_DISABLED_BY_DEFAULT);

BASE_FEATURE(kFormsClassificationsMqlsLogging,
             base::FEATURE_ENABLED_BY_DEFAULT);

BASE_FEATURE(kPasswordChangeSubmissionMqlsLogging,
             base::FEATURE_ENABLED_BY_DEFAULT);

BASE_FEATURE(kNotificationContentDetectionMqlsLogging,
             base::FEATURE_ENABLED_BY_DEFAULT);

BASE_FEATURE(kBlingPrototypingMqlsLogging, base::FEATURE_DISABLED_BY_DEFAULT);

BASE_FEATURE(kContextualTasksContextMqlsLogging,
             base::FEATURE_DISABLED_BY_DEFAULT);

BASE_FEATURE(kGeminiAntiscamProtectionMqlsLogging,
             base::FEATURE_DISABLED_BY_DEFAULT);

BASE_FEATURE(kFindsMqlsLogging, base::FEATURE_DISABLED_BY_DEFAULT);

BASE_FEATURE(kAtMemoryMqlsLogging, base::FEATURE_DISABLED_BY_DEFAULT);

}  // namespace features

namespace {

// Helper function that creates a `UserFeedbackCallback` for unspecified
// feedback.
UserFeedbackCallback FeedbackUnspecified() {
  return base::BindRepeating([](proto::LogAiDataRequest&) {
    return proto::UserFeedback::USER_FEEDBACK_UNSPECIFIED;
  });
}

void RegisterActorLogin() {
  MqlsFeatureRegistry::GetInstance().Register(
      std::make_unique<MqlsFeatureMetadata>(
          "ActorLogin", proto::LogAiDataRequest::FeatureCase::kActorLogin,
          /*enterprise_policy=*/std::nullopt, &features::kActorLoginMqlsLogging,
          FeedbackUnspecified()));
}

void RegisterCompose() {
  const char kComposeName[] = "Compose";
  EnterprisePolicyPref enterprise_policy =
      EnterprisePolicyRegistry::GetInstance().Register(
          prefs::kComposeEnterprisePolicyAllowed);

  UserFeedbackCallback logging_callback =
      base::BindRepeating([](proto::LogAiDataRequest& request_proto) {
        return request_proto.compose().quality().user_feedback();
      });
  auto mqls_metadata = std::make_unique<MqlsFeatureMetadata>(
      kComposeName, proto::LogAiDataRequest::FeatureCase::kCompose,
      enterprise_policy, &features::kComposeMqlsLogging, logging_callback);
  MqlsFeatureRegistry::GetInstance().Register(std::move(mqls_metadata));

  auto ui_metadata = std::make_unique<SettingsUiMetadata>(
      kComposeName, UserVisibleFeatureKey::kCompose, enterprise_policy);
  SettingsUiRegistry::GetInstance().Register(std::move(ui_metadata));
}

void RegisterWallpaperSearch() {
  const char kWallpaperSearchName[] = "WallpaperSearch";
  EnterprisePolicyPref enterprise_policy =
      EnterprisePolicyRegistry::GetInstance().Register(
          prefs::kWallpaperSearchEnterprisePolicyAllowed);

  UserFeedbackCallback logging_callback =
      base::BindRepeating([](proto::LogAiDataRequest& request_proto) {
        return request_proto.wallpaper_search().quality().user_feedback();
      });
  auto mqls_metadata = std::make_unique<MqlsFeatureMetadata>(
      kWallpaperSearchName,
      proto::LogAiDataRequest::FeatureCase::kWallpaperSearch, enterprise_policy,
      &features::kWallpaperSearchMqlsLogging, logging_callback);
  MqlsFeatureRegistry::GetInstance().Register(std::move(mqls_metadata));

  auto ui_metadata = std::make_unique<SettingsUiMetadata>(
      kWallpaperSearchName, UserVisibleFeatureKey::kWallpaperSearch,
      enterprise_policy);
  SettingsUiRegistry::GetInstance().Register(std::move(ui_metadata));
}

void RegisterHistorySearch() {
  EnterprisePolicyPref enterprise_policy =
      EnterprisePolicyRegistry::GetInstance().Register(
          prefs::kHistorySearchEnterprisePolicyAllowed);

  UserFeedbackCallback logging_callback_query =
      base::BindRepeating([](proto::LogAiDataRequest& request_proto) {
        return request_proto.history_query().quality().user_feedback();
      });
  auto mqls_metadata_query = std::make_unique<MqlsFeatureMetadata>(
      "HistoryQuery", proto::LogAiDataRequest::FeatureCase::kHistoryQuery,
      enterprise_policy, &features::kHistorySearchMqlsLogging,
      logging_callback_query);
  MqlsFeatureRegistry::GetInstance().Register(std::move(mqls_metadata_query));

  auto mqls_metadata_answer = std::make_unique<MqlsFeatureMetadata>(
      "HistoryAnswer", proto::LogAiDataRequest::FeatureCase::kHistoryAnswer,
      enterprise_policy, &features::kHistorySearchMqlsLogging,
      FeedbackUnspecified());
  MqlsFeatureRegistry::GetInstance().Register(std::move(mqls_metadata_answer));

  auto ui_metadata = std::make_unique<SettingsUiMetadata>(
      "HistorySearch", UserVisibleFeatureKey::kHistorySearch,
      enterprise_policy);
  SettingsUiRegistry::GetInstance().Register(std::move(ui_metadata));
}

void RegisterPasswordChangeSubmission() {
  const char kPasswordChangeSubmissionName[] = "PasswordChangeSubmission";
  EnterprisePolicyPref enterprise_policy =
      EnterprisePolicyRegistry::GetInstance().Register(
          prefs::kAutomatedPasswordChangeEnterprisePolicyAllowed);

  auto ui_metadata = std::make_unique<SettingsUiMetadata>(
      "PasswordChangeSubmission",
      UserVisibleFeatureKey::kPasswordChangeSubmission, enterprise_policy);
  SettingsUiRegistry::GetInstance().Register(std::move(ui_metadata));

  auto mqls_metadata = std::make_unique<MqlsFeatureMetadata>(
      kPasswordChangeSubmissionName,
      proto::LogAiDataRequest::FeatureCase::kPasswordChangeSubmission,
      enterprise_policy, &features::kPasswordChangeSubmissionMqlsLogging,
      FeedbackUnspecified());
  MqlsFeatureRegistry::GetInstance().Register(std::move(mqls_metadata));
}

void RegisterProductSpecifications() {
  EnterprisePolicyPref enterprise_policy =
      EnterprisePolicyRegistry::GetInstance().Register(
          prefs::kProductSpecificationsEnterprisePolicyAllowed);
  UserFeedbackCallback logging_callback =
      base::BindRepeating([](proto::LogAiDataRequest& request_proto) {
        return request_proto.product_specifications().quality().user_feedback();
      });
  auto metadata = std::make_unique<MqlsFeatureMetadata>(
      "ProductSpecifications",
      proto::LogAiDataRequest::FeatureCase::kProductSpecifications,
      enterprise_policy, &features::kProductSpecificationsMqlsLogging,
      logging_callback);
  MqlsFeatureRegistry::GetInstance().Register(std::move(metadata));
}

void RegisterAutofillPredictions() {
  MqlsFeatureRegistry::GetInstance().Register(
      std::make_unique<MqlsFeatureMetadata>(
          "FormsClassifications",
          proto::LogAiDataRequest::FeatureCase::kFormsClassifications,
          EnterprisePolicyRegistry::GetInstance().Register(
              prefs::kAutofillPredictionImprovementsEnterprisePolicyAllowed),
          &features::kFormsClassificationsMqlsLogging, FeedbackUnspecified()));
}

void RegisterNotificationContentDetection() {
  EnterprisePolicyPref enterprise_policy =
      EnterprisePolicyRegistry::GetInstance().Register(
          prefs::kNotificationContentDetectionEnterprisePolicyAllowed);
  UserFeedbackCallback logging_callback =
      base::BindRepeating([](proto::LogAiDataRequest& request_proto) {
        return request_proto.notification_content_detection()
            .quality()
            .user_feedback();
      });
  auto metadata = std::make_unique<MqlsFeatureMetadata>(
      "NotificationContentDetection",
      proto::LogAiDataRequest::FeatureCase::kNotificationContentDetection,
      enterprise_policy, &features::kNotificationContentDetectionMqlsLogging,
      logging_callback);
  MqlsFeatureRegistry::GetInstance().Register(std::move(metadata));
}

void RegisterBlingPrototyping() {
  MqlsFeatureRegistry::GetInstance().Register(
      std::make_unique<MqlsFeatureMetadata>(
          "BlingPrototyping",
          proto::LogAiDataRequest::FeatureCase::kBlingPrototyping,
          EnterprisePolicyRegistry::GetInstance().Register(
              prefs::kBlingPrototypingEnterprisePolicyAllowed),
          &features::kBlingPrototypingMqlsLogging, FeedbackUnspecified()));
}

void RegisterContextualTasksContext() {
  MqlsFeatureRegistry::GetInstance().Register(
      std::make_unique<MqlsFeatureMetadata>(
          "ContextualTasksContext",
          proto::LogAiDataRequest::FeatureCase::kContextualTasksContext,
          EnterprisePolicyRegistry::GetInstance().Register(
              prefs::kContextualTasksContextEnterprisePolicyAllowed),
          &features::kContextualTasksContextMqlsLogging,
          FeedbackUnspecified()));
}

void RegisterGeminiAntiscamProtection() {
  MqlsFeatureRegistry::GetInstance().Register(
      std::make_unique<MqlsFeatureMetadata>(
          "GeminiAntiscamProtection",
          proto::LogAiDataRequest::FeatureCase::kGeminiAntiscamProtection,
          EnterprisePolicyRegistry::GetInstance().Register(
              prefs::kGeminiAntiscamProtectionEnterprisePolicyAllowed),
          &features::kGeminiAntiscamProtectionMqlsLogging,
          FeedbackUnspecified()));
}

void RegisterFinds() {
  const char kFindsName[] = "Finds";
  EnterprisePolicyPref enterprise_policy =
      EnterprisePolicyRegistry::GetInstance().Register(
          prefs::kFindsEnterprisePolicyAllowed);

  auto ui_metadata = std::make_unique<SettingsUiMetadata>(
      kFindsName, UserVisibleFeatureKey::kFinds, enterprise_policy);
  SettingsUiRegistry::GetInstance().Register(std::move(ui_metadata));

  auto mqls_metadata = std::make_unique<MqlsFeatureMetadata>(
      kFindsName, proto::LogAiDataRequest::FeatureCase::kFinds,
      enterprise_policy, &features::kFindsMqlsLogging, FeedbackUnspecified());
  MqlsFeatureRegistry::GetInstance().Register(std::move(mqls_metadata));
}

void RegisterContextualCueing() {
  const char kContextualCueingName[] = "ContextualCueing";

  auto ui_metadata = std::make_unique<SettingsUiMetadata>(
      kContextualCueingName, UserVisibleFeatureKey::kContextualCueing,
      EnterprisePolicyRegistry::GetInstance().Register(
          prefs::kChromeSuggestionsSettings));
  SettingsUiRegistry::GetInstance().Register(std::move(ui_metadata));
}

void RegisterAtMemory() {
  const char kAtMemoryName[] = "AtMemory";

  EnterprisePolicyPref enterprise_policy =
      EnterprisePolicyRegistry::GetInstance().Register(
          prefs::kFindAndFillWithGeminiSettings);

  auto mqls_metadata = std::make_unique<MqlsFeatureMetadata>(
      kAtMemoryName, proto::LogAiDataRequest::FeatureCase::kAtMemory,
      enterprise_policy, &features::kAtMemoryMqlsLogging,
      FeedbackUnspecified());
  MqlsFeatureRegistry::GetInstance().Register(std::move(mqls_metadata));
}

}  // anonymous namespace

void RegisterGenAiFeatures(PrefRegistrySimple* pref_registry) {
  static bool features_registered = false;
  // When adding a value here, also update:
  // - tools/metrics/histograms/metadata/optimization_guide/histogram.xml:
  // <variants name="LogAiDataRequestFeature">
  if (!features_registered) {
    // The registries are static and so should only be populated once for the
    // program (rather than once per profile).
    RegisterActorLogin();
    RegisterCompose();
    RegisterWallpaperSearch();
    RegisterHistorySearch();
    RegisterProductSpecifications();
    RegisterAutofillPredictions();
    RegisterPasswordChangeSubmission();
    RegisterNotificationContentDetection();
    RegisterBlingPrototyping();
    RegisterContextualTasksContext();
    RegisterGeminiAntiscamProtection();
    RegisterFinds();
    RegisterContextualCueing();
    RegisterAtMemory();
    features_registered = true;
  }
  EnterprisePolicyRegistry::GetInstance().RegisterProfilePrefs(pref_registry);
}

}  // namespace optimization_guide