// 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.

#include "components/optimization_guide/core/model_execution/on_device_execution.h"

#include "base/metrics/histogram_functions.h"
#include "base/notimplemented.h"
#include "base/strings/strcat.h"
#include "base/strings/stringprintf.h"
#include "base/strings/to_string.h"
#include "base/trace_event/trace_event.h"
#include "components/optimization_guide/core/model_execution/model_execution_util.h"
#include "components/optimization_guide/core/model_execution/multimodal_message.h"
#include "components/optimization_guide/core/model_execution/on_device_features.h"
#include "components/optimization_guide/core/model_execution/repetition_checker.h"
#include "components/optimization_guide/core/optimization_guide_common.mojom.h"
#include "components/optimization_guide/core/optimization_guide_features.h"
#include "components/optimization_guide/public/mojom/model_broker.mojom.h"
#include "services/on_device_model/public/mojom/on_device_model.mojom.h"

namespace optimization_guide {

namespace {

using google::protobuf::RepeatedPtrField;

void LogRequest(OptimizationGuideLogger* logger,
                const proto::OnDeviceModelServiceRequest& logged_request) {
  if (logger && logger->ShouldEnableDebugLogs()) {
    OPTIMIZATION_GUIDE_LOGGER(
        optimization_guide_common::mojom::LogSource::MODEL_EXECUTION, logger)
        << "Executing model "
        << (logged_request.input_context_string().empty()
                ? ""
                : base::StringPrintf(
                      "with input context of %d tokens:\n%s\n",
                      logged_request.input_context_num_tokens_processed(),
                      logged_request.input_context_string()))
        << "with string:\n"
        << logged_request.execution_string();
  }
}

void LogRawResponse(OptimizationGuideLogger* logger,
                    mojom::OnDeviceFeature feature,
                    const std::string& raw_response) {
  if (logger && logger->ShouldEnableDebugLogs()) {
    OPTIMIZATION_GUIDE_LOGGER(
        optimization_guide_common::mojom::LogSource::MODEL_EXECUTION, logger)
        << "Model generates raw response with " << base::ToString(feature)
        << ":\n"
        << raw_response;
  }
}

void LogRepeatedResponse(OptimizationGuideLogger* logger,
                         mojom::OnDeviceFeature feature,
                         const std::string& repeated_response) {
  if (logger && logger->ShouldEnableDebugLogs()) {
    OPTIMIZATION_GUIDE_LOGGER(
        optimization_guide_common::mojom::LogSource::MODEL_EXECUTION, logger)
        << "Model generates repeated response with " << base::ToString(feature)
        << ":\n"
        << repeated_response;
  }
}

void LogResponseHasRepeats(mojom::OnDeviceFeature feature, bool has_repeats) {
  base::UmaHistogramBoolean(
      base::StrCat(
          {"OptimizationGuide.ModelExecution.OnDeviceResponseHasRepeats.",
           GetVariantName(feature)}),
      has_repeats);
}

std::string GenerateExecutionId() {
  return "on-device:" + base::Uuid::GenerateRandomV4().AsLowercaseString();
}

// Returns whether the feature tracks repetition.
// TODO(crbug.com/512149280): Move repetition checker to manifest config.
bool IsRepetitionTrackedFeature(mojom::OnDeviceFeature feature) {
  switch (feature) {
    case mojom::OnDeviceFeature::kProofreaderApi:
      return false;
    default:
      return true;
  }
}

}  // namespace

OnDeviceExecution::OnDeviceExecution(
    mojom::OnDeviceFeature feature,
    OnDeviceOptions opts,
    MultimodalMessage message,
    on_device_model::mojom::ResponseConstraintPtr constraint,
    std::unique_ptr<ResultLogger> logger,
    OptimizationGuideModelExecutionResultStreamingCallback callback,
    base::OnceClosure cleanup_callback)
    : feature_(feature),
      opts_(std::move(opts)),
      last_message_(std::move(message)),
      constraint_(std::move(constraint)),
      telemetry_logger_(feature),
      histogram_logger_(std::move(logger)),
      callback_(std::move(callback)),
      cleanup_callback_(std::move(cleanup_callback)) {
  exec_log_.set_execution_id(GenerateExecutionId());
  exec_log_.mutable_on_device_model_execution_info()->add_execution_infos();
  *(exec_log_.mutable_on_device_model_execution_info()
        ->mutable_model_versions()) = opts_.model_versions;
  // Note: if on-device fails for some reason, the result will be changed.
  histogram_logger_->set_result(Result::kUsedOnDevice);
}

OnDeviceExecution::~OnDeviceExecution() {
  if (callback_) {
    if (histogram_logger_) {
      histogram_logger_->set_result(Result::kDestroyedWhileWaitingForResponse);
    }
    telemetry_logger_.RecordDestroyedWhileWaiting();
  }
}

proto::OnDeviceModelServiceRequest* OnDeviceExecution::MutableLoggedRequest() {
  CHECK_GT(exec_log_.on_device_model_execution_info().execution_infos_size(),
           0);
  return exec_log_.mutable_on_device_model_execution_info()
      ->mutable_execution_infos(0)
      ->mutable_request()
      ->mutable_on_device_model_service_request();
}

proto::OnDeviceModelServiceResponse*
OnDeviceExecution::MutableLoggedResponse() {
  CHECK_GT(exec_log_.on_device_model_execution_info().execution_infos_size(),
           0);
  return exec_log_.mutable_on_device_model_execution_info()
      ->mutable_execution_infos(0)
      ->mutable_response()
      ->mutable_on_device_model_service_response();
}

void OnDeviceExecution::AddModelExecutionLogs(
    google::protobuf::RepeatedPtrField<
        proto::InternalOnDeviceModelExecutionInfo> logs) {
  exec_log_.mutable_on_device_model_execution_info()
      ->mutable_execution_infos()
      ->MergeFrom(std::move(logs));
}

void OnDeviceExecution::Cancel() {
  CancelPendingResponse(Result::kCancelled);
}

void OnDeviceExecution::BeginExecution(OnDeviceContext& context) {
  TRACE_EVENT("optimization_guide", "OnDeviceExecution::BeginExecution",
              "feature", base::ToString(feature_));
  auto input = opts_.adapter->ConstructInputString(
      last_message_.read(), /*want_input_context=*/false);
  if (!input) {
    CancelPendingResponse(Result::kFailedConstructingMessage);
    return;
  }

  auto* logged_request = MutableLoggedRequest();

  // Terminate optional context processing and log the context info.
  context.CloneSession(session_.BindNewPipeAndPassReceiver(), logged_request,
                       input->should_ignore_input_context);

  logged_request->set_execution_string(input->ToString());
  LogRequest(opts_.logger.get(), *logged_request);

  auto options = on_device_model::mojom::GenerateOptions::New();
  options->max_output_tokens = opts_.token_limits.max_output_tokens;
  options->constraint = constraint_ ? std::move(constraint_)
                                    : opts_.adapter->GetResponseConstraint();

  if (!options->constraint.is_null()) {
    auto hint_options = on_device_model::mojom::HintOptions::New();
    hint_options->constrained_decoding_hint = true;
    session_->Hint(std::move(hint_options));
  }

  if (input->input->pieces.size() > 0) {
    auto append_options = on_device_model::mojom::AppendOptions::New();
    append_options->input = std::move(input->input);
    append_options->max_tokens = opts_.token_limits.max_execute_tokens;
    session_->Append(std::move(append_options),
                     context_receiver_.BindNewPipeAndPassRemote());
  }

  opts_.safety_checker->RunRequestChecks(
      last_message_,
      base::BindOnce(&OnDeviceExecution::OnRequestSafetyResult,
                     weak_ptr_factory_.GetWeakPtr(), std::move(options)));
}

void OnDeviceExecution::OnRequestSafetyResult(
    on_device_model::mojom::GenerateOptionsPtr options,
    SafetyChecker::Result safety_result) {
  TRACE_EVENT("optimization_guide", "OnDeviceExecution::OnRequestSafetyResult",
              "feature", base::ToString(feature_));
  if (safety_result.failed_to_run) {
    CancelPendingResponse(Result::kFailedConstructingMessage,
                          OnDeviceError::kFailedToRunSafety);
    return;
  }
  // Log the check executions.
  AddModelExecutionLogs(std::move(safety_result.logs));

  // Handle the result.
  if (safety_result.is_unsafe || safety_result.is_unsupported_language) {
    CancelPendingResponse(Result::kRequestUnsafe,
                          safety_result.is_unsupported_language
                              ? OnDeviceError::kUnsupportedLanguage
                              : OnDeviceError::kFiltered);
    return;
  }
  BeginRequestExecution(std::move(options));
}

void OnDeviceExecution::BeginRequestExecution(
    on_device_model::mojom::GenerateOptionsPtr options) {
  session_->Generate(std::move(options), receiver_.BindNewPipeAndPassRemote());
  receiver_.set_disconnect_with_reason_handler(base::BindOnce(
      &OnDeviceExecution::OnResponderDisconnect, base::Unretained(this)));
}

// on_device_model::mojom::StreamingResponder:
void OnDeviceExecution::OnResponse(
    on_device_model::mojom::ResponseChunkPtr chunk) {
  TRACE_EVENT("optimization_guide.debug", "OnDeviceExecution::OnResponse",
              "feature", base::ToString(feature_));
  proto::OnDeviceModelServiceResponse* logged_response =
      MutableLoggedResponse();

  if (current_response_.empty()) {
    telemetry_logger_.RecordFirstResponse();
    logged_response->set_time_to_first_response_millis(
        telemetry_logger_.GetTimeToFirstResponse().InMilliseconds());
  }

  NewlineBuffer::Chunk trimmed_chunk = newline_buffer_.Append(chunk->text);
  if (trimmed_chunk.text.empty()) {
    return;
  }
  current_response_ += trimmed_chunk.text;
  num_unchecked_response_tokens_ += trimmed_chunk.num_tokens;
  num_response_tokens_ += trimmed_chunk.num_tokens;

  if (IsRepetitionTrackedFeature(feature_) &&
      HasRepeatingSuffix(current_response_)) {
    // If a repeat is detected, halt the response, and cancel/finish early.
    receiver_.reset();
    logged_response->set_has_repeats(true);
    if (features::GetOnDeviceModelRetractRepeats()) {
      LogRepeatedResponse(opts_.logger.get(), feature_, current_response_);
      logged_response->set_status(
          proto::ON_DEVICE_MODEL_SERVICE_RESPONSE_STATUS_RETRACTED);
      CancelPendingResponse(Result::kResponseHadRepeats,
                            OnDeviceError::kResponseLowQuality);
      return;
    }

    // Artificially send the OnComplete event to finish processing.
    OnComplete(on_device_model::mojom::ResponseSummary::New());
    return;
  }

  if (!opts_.safety_checker->safety_cfg().CanCheckPartialOutput(
          num_response_tokens_, num_unchecked_response_tokens_)) {
    // Not enough new data to be worth re-evaluating yet.
    return;
  }

  num_unchecked_response_tokens_ = 0;
  RunRawOutputSafetyCheck(ResponseCompleteness::kPartial);
}

void OnDeviceExecution::OnComplete(
    on_device_model::mojom::ResponseSummaryPtr summary) {
  TRACE_EVENT("optimization_guide", "OnDeviceExecution::OnComplete", "feature",
              base::ToString(feature_));
  receiver_.reset();  // Suppress expected disconnect

  bool has_repeats = MutableLoggedResponse()->has_repeats();

  LogResponseHasRepeats(feature_, has_repeats);
  telemetry_logger_.RecordCompletion(num_response_tokens_);
  MutableLoggedResponse()->set_time_to_completion_millis(
      telemetry_logger_.GetTimeToCompletion().InMilliseconds());

  output_token_count_ = summary->output_token_count;

  opts_.model_client->OnResponseCompleted();

  RunRawOutputSafetyCheck(ResponseCompleteness::kComplete);
}

void OnDeviceExecution::OnToolCalls(
    std::vector<on_device_model::mojom::ToolCallPtr> tool_calls) {
  // Tool calls are unexpected in the optimization guide execution path since
  // it never declares tools. Report as a bad message from the backend.
  receiver_.ReportBadMessage(
      "Unexpected tool calls in optimization guide execution path.");
  CancelPendingResponse(Result::kDisconnectAndCancel,
                        OnDeviceError::kGenericFailure);
}

void OnDeviceExecution::OnComplete(uint32_t tokens_processed) {
  TRACE_EVENT("optimization_guide",
              "OnDeviceExecution::[ContextClient]::OnComplete", "feature",
              base::ToString(feature_));
  execute_input_token_count_ = tokens_processed;
  MutableLoggedRequest()->set_execution_num_tokens_processed(tokens_processed);
}

void OnDeviceExecution::OnResponderDisconnect(uint32_t custom_reason,
                                              const std::string& description) {
  TRACE_EVENT("optimization_guide", "OnDeviceExecution::OnResponse", "feature",
              base::ToString(feature_));
  // OnComplete resets the receiver, so this implies that the response is
  // incomplete and there was either a service crash, error, or model eviction.
  receiver_.reset();
  switch (static_cast<on_device_model::mojom::GenerateError>(custom_reason)) {
    case on_device_model::mojom::GenerateError::kUnknown:
      CancelPendingResponse(Result::kDisconnectAndCancel);
      break;
    case on_device_model::mojom::GenerateError::kInvalidConstraint:
      CancelPendingResponse(Result::kFailedConstructingMessage,
                            OnDeviceError::kInvalidRequest);
      break;
  }
}

void OnDeviceExecution::RunRawOutputSafetyCheck(
    ResponseCompleteness completeness) {
  opts_.safety_checker->RunRawOutputCheck(
      current_response_, completeness,
      base::BindOnce(&OnDeviceExecution::OnRawOutputSafetyResult,
                     weak_ptr_factory_.GetWeakPtr(), current_response_.size(),
                     completeness));
}

void OnDeviceExecution::OnRawOutputSafetyResult(
    size_t raw_output_size,
    ResponseCompleteness completeness,
    SafetyChecker::Result safety_result) {
  TRACE_EVENT("optimization_guide.debug",
              "OnDeviceExecution::OnRawOutputSafetyResult", "feature",
              base::ToString(feature_));
  if (safety_result.failed_to_run) {
    CancelPendingResponse(Result::kFailedConstructingMessage,
                          OnDeviceError::kFailedToRunSafety);
    return;
  }
  if (safety_result.is_unsafe || safety_result.is_unsupported_language) {
    if (opts_.safety_checker->safety_cfg()
            .OnlyCancelUnsafeResponseOnComplete() &&
        completeness != ResponseCompleteness::kComplete) {
      return;
    }
    AddModelExecutionLogs(std::move(safety_result.logs));
    CancelPendingResponse(Result::kUsedOnDeviceOutputUnsafe,
                          safety_result.is_unsupported_language
                              ? OnDeviceError::kUnsupportedLanguage
                              : OnDeviceError::kFiltered);
    return;
  }
  if (completeness == ResponseCompleteness::kComplete) {
    AddModelExecutionLogs(std::move(safety_result.logs));
  }
  latest_safe_raw_output_.length = raw_output_size;
  MaybeParseResponse(completeness);
}

void OnDeviceExecution::MaybeParseResponse(ResponseCompleteness completeness) {
  if (!opts_.adapter->ShouldParseResponse(completeness)) {
    return;
  }

  std::string safe_response =
      current_response_.substr(0, latest_safe_raw_output_.length);
  LogRawResponse(opts_.logger.get(), feature_, safe_response);
  MutableLoggedResponse()->set_output_string(safe_response);
  size_t previous_response_pos = latest_response_pos_;
  latest_response_pos_ = latest_safe_raw_output_.length;
  opts_.adapter->ParseResponse(
      last_message_, safe_response, previous_response_pos,
      base::BindOnce(&OnDeviceExecution::OnParsedResponse,
                     weak_ptr_factory_.GetWeakPtr(), completeness));
}

void OnDeviceExecution::OnParsedResponse(
    ResponseCompleteness completeness,
    base::expected<proto::Any, ResponseParsingError> output) {
  TRACE_EVENT("optimization_guide.debug", "OnDeviceExecution::OnParsedResponse",
              "feature", base::ToString(feature_));
  if (!output.has_value()) {
    switch (output.error()) {
      case ResponseParsingError::kRejectedPii:
        MutableLoggedResponse()->set_status(
            proto::ON_DEVICE_MODEL_SERVICE_RESPONSE_STATUS_RETRACTED);
        CancelPendingResponse(Result::kContainedPII, OnDeviceError::kFiltered);
        return;
      case ResponseParsingError::kInvalidConfiguration:
      case ResponseParsingError::kFailed:
        CancelPendingResponse(Result::kFailedConstructingResponseMessage,
                              OnDeviceError::kResponseParsingFailed);
        return;
    }
  }
  opts_.safety_checker->RunResponseChecks(
      last_message_, *output, completeness,
      base::BindOnce(&OnDeviceExecution::OnResponseSafetyResult,
                     weak_ptr_factory_.GetWeakPtr(), completeness, *output));
}

void OnDeviceExecution::OnResponseSafetyResult(
    ResponseCompleteness completeness,
    proto::Any output,
    SafetyChecker::Result safety_result) {
  TRACE_EVENT("optimization_guide.debug",
              "OnDeviceExecution::OnResponseSafetyResult", "feature",
              base::ToString(feature_));
  if (safety_result.failed_to_run) {
    CancelPendingResponse(Result::kFailedConstructingMessage,
                          OnDeviceError::kFailedToRunSafety);
    return;
  }
  if (completeness == ResponseCompleteness::kComplete ||
      safety_result.is_unsafe || safety_result.is_unsupported_language) {
    AddModelExecutionLogs(std::move(safety_result.logs));
  }
  if (safety_result.is_unsafe || safety_result.is_unsupported_language) {
    if (opts_.safety_checker->safety_cfg()
            .OnlyCancelUnsafeResponseOnComplete() &&
        completeness != ResponseCompleteness::kComplete) {
      return;
    }
    CancelPendingResponse(Result::kUsedOnDeviceOutputUnsafe,
                          safety_result.is_unsupported_language
                              ? OnDeviceError::kUnsupportedLanguage
                              : OnDeviceError::kFiltered);
    return;
  }
  if (completeness == ResponseCompleteness::kPartial) {
    SendPartialResponseCallback(output);
    return;
  }

  SendSuccessCompletionCallback(output);
}

void OnDeviceExecution::CancelPendingResponse(Result result,
                                              OnDeviceError error) {
  TRACE_EVENT("optimization_guide", "OnDeviceExecution::CancelPendingResponse",
              "feature", base::ToString(feature_));
  if (!callback_) {
    return;
  }
  if (histogram_logger_) {
    histogram_logger_->set_result(result);
  }
  exec_log_.set_model_execution_error_enum(static_cast<uint32_t>(error));
  auto self = weak_ptr_factory_.GetWeakPtr();
  std::move(callback_).Run(OptimizationGuideModelStreamingExecutionResult(
      base::unexpected(error), /*provided_by_on_device=*/true,
      std::make_unique<proto::ModelExecutionInfo>(std::move(exec_log_))));
  if (self) {
    self->Cleanup();
  }
}

void OnDeviceExecution::SendPartialResponseCallback(
    const proto::Any& success_response_metadata) {
  TRACE_EVENT("optimization_guide.debug",
              "OnDeviceExecution::SendPartialResponseCallback", "feature",
              base::ToString(feature_));
  callback_.Run(OptimizationGuideModelStreamingExecutionResult(
      base::ok(StreamingResponse{.response = success_response_metadata,
                                 .is_complete = false}),
      /*provided_by_on_device=*/true));
}

void OnDeviceExecution::SendSuccessCompletionCallback(
    const proto::Any& success_response_metadata) {
  TRACE_EVENT("optimization_guide",
              "OnDeviceExecution::SendSuccessCompletionCallback", "feature",
              base::ToString(feature_));
  MutableLoggedResponse()->set_status(
      proto::ON_DEVICE_MODEL_SERVICE_RESPONSE_STATUS_SUCCESS);
  // Return the execution response.
  auto self = weak_ptr_factory_.GetWeakPtr();
  std::move(callback_).Run(OptimizationGuideModelStreamingExecutionResult(
      base::ok(
          StreamingResponse{.response = success_response_metadata,
                            .is_complete = true,
                            .input_token_count = execute_input_token_count_,
                            .output_token_count = output_token_count_}),
      /*provided_by_on_device=*/true,
      std::make_unique<proto::ModelExecutionInfo>(std::move(exec_log_))));
  if (self) {
    self->Cleanup();
  }
}

void OnDeviceExecution::Cleanup() {
  weak_ptr_factory_.InvalidateWeakPtrs();
  session_.reset();
  receiver_.reset();
  context_receiver_.reset();
  callback_.Reset();
  exec_log_.Clear();
  current_response_.clear();
  histogram_logger_.reset();
  std::move(cleanup_callback_).Run();
}

OnDeviceExecution::SafeRawOutput::SafeRawOutput() = default;
OnDeviceExecution::SafeRawOutput::~SafeRawOutput() = default;

OnDeviceExecution::ResultLogger::~ResultLogger() {
  base::UmaHistogramEnumeration(
      base::StrCat(
          {"OptimizationGuide.ModelExecution.OnDeviceExecuteModelResult.",
           GetVariantName(feature_)}),
      result_);
}

}  // namespace optimization_guide