// 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 SERVICES_WEBNN_TFLITE_CONTEXT_IMPL_LITERT_H_
#define SERVICES_WEBNN_TFLITE_CONTEXT_IMPL_LITERT_H_

#include <optional>

#include "mojo/public/cpp/bindings/pending_associated_receiver.h"
#include "services/webnn/public/cpp/webnn_trace.h"
#include "services/webnn/public/cpp/webnn_types.h"
#include "services/webnn/webnn_context_impl.h"
#include "services/webnn/webnn_graph_impl.h"

namespace webnn {

class WebNNConstantOperand;
class WebNNContextProviderInRenderer;

namespace litert {

// `ContextImplLiteRt` is created by `WebNNContextProviderImpl` and responsible
// for creating a `GraphImplLiteRt` which uses LiteRt for inference.
class ContextImplLiteRt final : public WebNNContextImpl {
 public:
  // Constructs a new `ContextImplLiteRt`. Must be called on
  // `owning_task_runner`.
  static std::unique_ptr<WebNNContextImpl, OnTaskRunnerDeleter> Create(
      mojo::PendingReceiver<mojom::WebNNContext> receiver,
      base::WeakPtr<WebNNContextProviderImpl> context_provider,
      mojom::CreateContextOptionsPtr options,
      mojo::ScopedDataPipeConsumerHandle write_tensor_consumer,
      mojo::ScopedDataPipeProducerHandle read_tensor_producer,
      std::unique_ptr<GpuTaskScheduler> gpu_task_scheduler,
      scoped_refptr<gpu::MemoryTracker> memory_tracker,
      scoped_refptr<base::SingleThreadTaskRunner> owning_task_runner,
      gpu::SharedImageManager* shared_image_manager,
      scoped_refptr<base::SingleThreadTaskRunner> main_task_runner,
      ScopedTrace scoped_trace,
      bool is_incognito);

  // Factory method for running without GPU dependencies (e.g., in the renderer
  // process).
  static WebNNContextImplPtr CreateForRenderer(
      mojo::PendingReceiver<mojom::WebNNContext> receiver,
      base::WeakPtr<WebNNContextProviderInRenderer> context_provider,
      mojom::CreateContextOptionsPtr options,
      scoped_refptr<base::SingleThreadTaskRunner> owning_task_runner,
      scoped_refptr<base::SingleThreadTaskRunner> main_task_runner);

  ContextImplLiteRt(
      mojo::PendingReceiver<mojom::WebNNContext> receiver,
      base::WeakPtr<WebNNContextProviderImpl> context_provider,
      mojom::CreateContextOptionsPtr options,
      mojo::ScopedDataPipeConsumerHandle write_tensor_consumer,
      mojo::ScopedDataPipeProducerHandle read_tensor_producer,
      std::unique_ptr<GpuTaskScheduler> gpu_task_scheduler,
      scoped_refptr<gpu::MemoryTracker> memory_tracker,
      scoped_refptr<base::SingleThreadTaskRunner> owning_task_runner,
      gpu::SharedImageManager* shared_image_manager,
      scoped_refptr<base::SingleThreadTaskRunner> main_task_runner,
      bool is_incognito);

  // Constructor for running without GPU dependencies.
  ContextImplLiteRt(
      mojo::PendingReceiver<mojom::WebNNContext> receiver,
      base::WeakPtr<WebNNContextProviderInRenderer> context_provider,
      mojom::CreateContextOptionsPtr options,
      scoped_refptr<base::SingleThreadTaskRunner> owning_task_runner,
      scoped_refptr<base::SingleThreadTaskRunner> main_task_runner);

  ContextImplLiteRt(const WebNNContextImpl&) = delete;
  ContextImplLiteRt& operator=(const ContextImplLiteRt&) = delete;

  // WebNNContextImpl:
  base::WeakPtr<WebNNContextImpl> AsWeakPtr() override;

 private:
  ~ContextImplLiteRt() override;

  void CreateGraphImpl(
      mojom::GraphInfoPtr graph_info,
      WebNNGraphImpl::ComputeResourceInfo compute_resource_info,
      base::flat_map<OperandId, std::unique_ptr<WebNNConstantOperand>>
          constant_operands,
      CreateGraphImplCallback callback) override;

  void DidOpenWeightsFile(
      mojom::GraphInfoPtr graph_info,
      WebNNGraphImpl::ComputeResourceInfo compute_resource_info,
      base::flat_map<OperandId, std::unique_ptr<WebNNConstantOperand>>
          constant_operands,
      CreateGraphImplCallback callback,
      base::File weights_file,
      mojo::PendingRemote<mojom::WeightsFileSession> session);

  void DidCreateWeightsFile(
      mojom::GraphInfoPtr graph_info,
      WebNNGraphImpl::ComputeResourceInfo compute_resource_info,
      base::flat_map<OperandId, std::unique_ptr<WebNNConstantOperand>>
          constant_operands,
      CreateGraphImplCallback callback,
      base::File weights_file);

  base::expected<scoped_refptr<WebNNTensorImpl>, mojom::ErrorPtr>
  CreateTensorImpl(mojo::PendingAssociatedReceiver<mojom::WebNNTensor> receiver,
                   mojom::TensorInfoPtr tensor_info) override;

  base::expected<scoped_refptr<WebNNTensorImpl>, mojom::ErrorPtr>
  CreateTensorFromSharedImageImpl(
      mojo::PendingAssociatedReceiver<mojom::WebNNTensor> receiver,
      mojom::TensorInfoPtr tensor_info,
      WebNNTensorImpl::RepresentationPtr representation) override;

  std::string_view GetBackendName() const override;
  std::vector<mojom::WebNNExecutionProviderDetailsPtr>
  GetExecutionProvidersInfo() const override;

  // Only be used in the GPU-process flow to indicate whether the profile is in
  // incognito.
  // For the LiteRT in renderer-process, the incognito mode flag will be
  // checked on the browser side to create temporary weight files or invalid
  // files.
  const std::optional<bool> is_incognito_;
  base::WeakPtrFactory<ContextImplLiteRt> weak_factory_{this};
};

}  // namespace litert
}  // namespace webnn

#endif  // SERVICES_WEBNN_TFLITE_CONTEXT_IMPL_LITERT_H_