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

#include "ml/text_classifier_impl.h"

#include <utility>
#include <vector>

#include <base/logging.h>
#include <utils/utf8/unicodetext.h>

#include "ml/mojom/text_classifier.mojom.h"
#include "ml/request_metrics.h"

namespace ml {

namespace {

using ::chromeos::machine_learning::mojom::CodepointSpan;
using ::chromeos::machine_learning::mojom::SuggestSelectionResult;
using ::chromeos::machine_learning::mojom::TextAnnotation;
using ::chromeos::machine_learning::mojom::TextAnnotationPtr;
using ::chromeos::machine_learning::mojom::TextAnnotationRequestPtr;
using ::chromeos::machine_learning::mojom::TextAnnotationResult;
using ::chromeos::machine_learning::mojom::TextClassifierRequest;
using ::chromeos::machine_learning::mojom::TextEntity;
using ::chromeos::machine_learning::mojom::TextEntityData;
using ::chromeos::machine_learning::mojom::TextEntityPtr;
using ::chromeos::machine_learning::mojom::TextSuggestSelectionRequestPtr;

// To avoid passing a lambda as a base::Closure.
void DeleteTextClassifierImpl(
    const TextClassifierImpl* const text_classifier_impl) {
  delete text_classifier_impl;
}

}  // namespace

bool TextClassifierImpl::Create(
    std::unique_ptr<libtextclassifier3::ScopedMmap>* mmap,
    TextClassifierRequest request) {
  auto text_classifier_impl = new TextClassifierImpl(mmap, std::move(request));
  if (text_classifier_impl->annotator_ == nullptr) {
    // Fails to create annotator, return nullptr.
    delete text_classifier_impl;
    return false;
  }

  // Use a connection error handler to strongly bind |text_classifier_impl| to
  // |request|.
  text_classifier_impl->SetConnectionErrorHandler(base::Bind(
      &DeleteTextClassifierImpl, base::Unretained(text_classifier_impl)));
  return true;
}

TextClassifierImpl::TextClassifierImpl(
    std::unique_ptr<libtextclassifier3::ScopedMmap>* mmap,
    TextClassifierRequest request)
    : annotator_(libtextclassifier3::Annotator::FromScopedMmap(
          mmap, nullptr, nullptr)),
      binding_(this, std::move(request)) {}

void TextClassifierImpl::SetConnectionErrorHandler(
    base::Closure connection_error_handler) {
  binding_.set_connection_error_handler(std::move(connection_error_handler));
}

void TextClassifierImpl::Annotate(TextAnnotationRequestPtr request,
                                  AnnotateCallback callback) {
  RequestMetrics<TextAnnotationResult> request_metrics("TextClassifier",
                                                       "Annotate");
  request_metrics.StartRecordingPerformanceMetrics();

  // Parse and set up the options.
  libtextclassifier3::AnnotationOptions option;
  if (request->default_locales) {
    option.locales = request->default_locales.value();
  }
  if (request->reference_time) {
    option.reference_time_ms_utc =
        request->reference_time->ToTimeT() * base::Time::kMillisecondsPerSecond;
  }
  if (request->reference_timezone) {
    option.reference_timezone = request->reference_timezone.value();
  }
  if (request->enabled_entities) {
    option.entity_types.insert(request->enabled_entities.value().begin(),
                               request->enabled_entities.value().end());
  }
  option.detected_text_language_tags =
      request->detected_text_language_tags.value_or("en");
  option.annotation_usecase =
      static_cast<libtextclassifier3::AnnotationUsecase>(
          request->annotation_usecase);

  // Do the annotation.
  const std::vector<libtextclassifier3::AnnotatedSpan> annotated_spans =
      annotator_->Annotate(request->text, option);

  // Parse the result.
  std::vector<TextAnnotationPtr> annotations;
  for (const auto& annotated_result : annotated_spans) {
    DCHECK(annotated_result.span.second >= annotated_result.span.first);
    std::vector<TextEntityPtr> entities;
    for (const auto& classification : annotated_result.classification) {
      // First, get entity data.
      auto entity_data = TextEntityData::New();
      if (classification.collection == "number") {
        entity_data->set_numeric_value(classification.numeric_double_value);
      } else {
        // For the other types, just encode the substring into string_value.
        // TODO(honglinyu): add data extraction for more types when needed
        // and available.
        // Note that the returned indices by annotator is unicode codepoints.
        entity_data->set_string_value(
            libtextclassifier3::UTF8ToUnicodeText(request->text, false)
                .UTF8Substring(annotated_result.span.first,
                               annotated_result.span.second));
      }

      // Second, create the entity.
      entities.emplace_back(TextEntity::New(classification.collection,
                                            classification.score,
                                            std::move(entity_data)));
    }
    annotations.emplace_back(TextAnnotation::New(annotated_result.span.first,
                                                 annotated_result.span.second,
                                                 std::move(entities)));
  }

  std::move(callback).Run(std::move(annotations));

  request_metrics.FinishRecordingPerformanceMetrics();
  request_metrics.RecordRequestEvent(TextAnnotationResult::OK);
}

void TextClassifierImpl::SuggestSelection(
    TextSuggestSelectionRequestPtr request, SuggestSelectionCallback callback) {
  RequestMetrics<SuggestSelectionResult> request_metrics("TextClassifier",
                                                         "SuggestSelection");
  request_metrics.StartRecordingPerformanceMetrics();

  libtextclassifier3::SelectionOptions option;
  if (request->default_locales) {
    option.locales = request->default_locales.value();
  }
  option.detected_text_language_tags =
      request->detected_text_language_tags.value_or("en");
  option.annotation_usecase =
      static_cast<libtextclassifier3::AnnotationUsecase>(
          request->annotation_usecase);

  libtextclassifier3::CodepointSpan user_selection;
  user_selection.first = request->user_selection->start_offset;
  user_selection.second = request->user_selection->end_offset;

  const libtextclassifier3::CodepointSpan suggested_span =
      annotator_->SuggestSelection(request->text, user_selection, option);
  auto result_span = CodepointSpan::New();
  result_span->start_offset = suggested_span.first;
  result_span->end_offset = suggested_span.second;

  std::move(callback).Run(std::move(result_span));

  request_metrics.FinishRecordingPerformanceMetrics();
  request_metrics.RecordRequestEvent(SuggestSelectionResult::OK);
}

}  // namespace ml