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

#ifndef COMPONENTS_LIVE_CAPTION_TRANSLATION_DISPATCHER_ON_DEVICE_H_
#define COMPONENTS_LIVE_CAPTION_TRANSLATION_DISPATCHER_ON_DEVICE_H_

#include <memory>
#include <optional>
#include <string>
#include <vector>

#include "base/memory/weak_ptr.h"
#include "components/live_caption/translation_dispatcher.h"
#include "components/on_device_translation/public/mojom/translator.mojom.h"
#include "components/on_device_translation/service_controller.h"
#include "mojo/public/cpp/bindings/remote.h"
#include "services/network/public/mojom/url_loader_factory.mojom.h"
#include "url/origin.h"

namespace captions {

// These values are persisted to logs. Entries should not be renumbered and
// numeric values should never be reused.
// LINT.IfChange(OnDeviceTranslationErrorReason)
enum class OnDeviceTranslationErrorReason {
  kCanTranslateLanguageNotSupported = 0,
  kCanTranslateServiceCrashed = 1,
  kCanTranslateExceedsServiceCountLimitation = 2,
  kCreateTranslatorInvalidBinary = 3,
  kCreateTranslatorInvalidFunctionPointer = 4,
  kCreateTranslatorFailedToInitialize = 5,
  kCreateTranslatorFailedToCreateTranslator = 6,
  kCreateTranslatorInvalidVersion = 7,
  kCreateTranslatorServiceCrashed = 8,
  kCreateTranslatorNotSupportedLanguage = 9,
  kCreateTranslatorExceedsServiceCountLimitation = 10,
  kCreateTranslatorExceedsPendingTaskCountLimitation = 11,
  kCreateTranslatorUnknownError = 12,
  kTranslationExecutionFailed = 13,
  kMaxValue = kTranslationExecutionFailed,
};
// LINT.ThenChange(//tools/metrics/histograms/metadata/accessibility/enums.xml:OnDeviceTranslationErrorReason)

class TranslationDispatcherOnDevice : public TranslationDispatcher {
 public:
  TranslationDispatcherOnDevice();
  ~TranslationDispatcherOnDevice() override;
  TranslationDispatcherOnDevice(const TranslationDispatcherOnDevice&) = delete;
  TranslationDispatcherOnDevice& operator=(
      const TranslationDispatcherOnDevice&) = delete;
  explicit TranslationDispatcherOnDevice(
      std::unique_ptr<on_device_translation::OnDeviceTranslationController>
          translation_controller);

  void GetTranslation(std::string_view result,
                      std::string_view source_language,
                      std::string_view target_language,
                      TranslateEventCallback callback) override;

  void SetURLLoaderFactoryForTest(
      mojo::Remote<network::mojom::URLLoaderFactory> url_loader_factory);

 private:
  std::string GetBaseLanguage(const std::string& locale_code);

  void OnTranslationCreated(
      const std::string& source_language,
      const std::string& target_language,
      const std::string& result,
      TranslateEventCallback callback,
      base::expected<
          mojo::PendingRemote<on_device_translation::mojom::OnDeviceTranslator>,
          on_device_translation::OnDeviceTranslationController::
              CreateTranslatorError> translator);

  void OnTranslated(TranslateEventCallback callback,
                    const std::optional<std::string>& translation);

  void OnCanTranslate(
      const std::string& source_language,
      const std::string& target_language,
      const std::string& result,
      TranslateEventCallback callback,
      on_device_translation::OnDeviceTranslationController::CanTranslateResult
          can_translate_result);

  void ResetCreationAndNotifyFailure(
      TranslateEventCallback callback,
      OnDeviceTranslationErrorReason error_reason);

  std::string source_language_;
  std::string target_language_;

  std::unique_ptr<on_device_translation::OnDeviceTranslationController>
      translation_controller_;
  mojo::Remote<on_device_translation::mojom::OnDeviceTranslator> translator_;
  bool creation_in_progress_ = false;
  std::vector<std::pair<std::string, TranslateEventCallback>>
      pending_callbacks_;

  base::WeakPtrFactory<TranslationDispatcherOnDevice> weak_factory_{this};
};

}  // namespace captions

#endif  // COMPONENTS_LIVE_CAPTION_TRANSLATION_DISPATCHER_ON_DEVICE_H_