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

#ifndef SERVICES_WEBNN_ORT_CONTEXT_IMPL_ORT_H_
#define SERVICES_WEBNN_ORT_CONTEXT_IMPL_ORT_H_

#include "base/memory/scoped_refptr.h"
#include "base/task/cancelable_task_tracker.h"
#include "mojo/public/cpp/bindings/pending_associated_receiver.h"
#include "services/webnn/ort/device_allocator.h"
#include "services/webnn/ort/environment.h"
#include "services/webnn/ort/ort_session_options.h"
#include "services/webnn/ort/scoped_ort_types.h"
#include "services/webnn/public/cpp/webnn_trace.h"
#include "services/webnn/public/cpp/webnn_types.h"
#include "services/webnn/webnn_context_impl.h"

namespace webnn {

class WebNNConstantOperand;

namespace ort {

// `ContextImplOrt` is created by `WebNNContextProviderImpl` and responsible
// for creating a `GraphImplOrt` which uses ONNX Runtime for inference.
class ContextImplOrt : public WebNNContextImpl {
 public:
  // Constructs a new `ContextImplOrt`. 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,
      scoped_refptr<Environment> env,
      scoped_refptr<SessionOptions> session_options,
      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);

  ContextImplOrt(mojo::PendingReceiver<mojom::WebNNContext> receiver,
                 base::WeakPtr<WebNNContextProviderImpl> context_provider,
                 const EpWorkarounds& ep_workarounds,
                 mojom::CreateContextOptionsPtr options,
                 scoped_refptr<SessionOptions> session_options,
                 mojo::ScopedDataPipeConsumerHandle write_tensor_consumer,
                 mojo::ScopedDataPipeProducerHandle read_tensor_producer,
                 scoped_refptr<Environment> env,
                 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);

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


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

  static ContextProperties GetContextProperties(
      bool resample2d_limit_to_nchw);

  scoped_refptr<Environment> env() const { return env_; }

  scoped_refptr<SessionOptions> session_options() const {
    return session_options_;
  }

  base::CancelableTaskTracker& cancelable_task_tracker() {
    return cancelable_task_tracker_;
  }

  void HandleContextLostOrCrash(const std::string& error_message,
                                OrtErrorCode error_code);

 protected:
  ~ContextImplOrt() override;

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

  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;

  scoped_refptr<Environment> env_;

  // The session options are shared among all the sessions created by this
  // context.
  scoped_refptr<SessionOptions> session_options_;

  // The device allocator used for device tensor creation. May be nullptr if
  // device tensor is not supported.
  scoped_refptr<DeviceAllocator> device_allocator_;

  // The importer for external resources like D3D12 buffers. It is used for
  // importing D3D12 buffers into ORT tensors. May be nullptr if the EP does
  // not support external resource import.
  ScopedOrtExternalResourceImporter external_resource_importer_;

  // Cancels pending graph compilation tasks when destructing.
  base::CancelableTaskTracker cancelable_task_tracker_;

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

}  // namespace ort
}  // namespace webnn

#endif  // SERVICES_WEBNN_ORT_CONTEXT_IMPL_ORT_H_