// Copyright 2024 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_WEBNN_TENSOR_IMPL_H_
#define SERVICES_WEBNN_WEBNN_TENSOR_IMPL_H_

#include "base/component_export.h"
#include "base/functional/bind.h"
#include "base/functional/callback_helpers.h"
#include "base/memory/raw_ref.h"
#include "gpu/command_buffer/common/sync_token.h"
#include "gpu/command_buffer/service/shared_image/shared_image_representation.h"
#include "mojo/public/cpp/base/big_buffer.h"
#include "mojo/public/cpp/bindings/associated_receiver.h"
#include "mojo/public/cpp/bindings/pending_associated_receiver.h"
#include "services/webnn/public/cpp/operand_descriptor.h"
#include "services/webnn/public/mojom/webnn_tensor.mojom.h"
#include "services/webnn/sequence_deleter.h"
#include "services/webnn/webnn_object_impl.h"

namespace webnn {

class WebNNContextImpl;

// GPU process implementation of the `MLTensor` interface. While this class is
// reference-counted a `WebNNTensorImpl` is guaranteed not to outlive the
// `WebNNContextImpl` that created it because references are only held by the
// context itself or by tasks scheduled to its `gpu::Scheduler` sequence which
// is shut down when the context is destroyed.
//
// This invariant is checked by the `raw_ref<WebNNContextImpl>` member, which
// will trigger dangling pointer warnings in debug builds and safe crashes in
// release builds.
class COMPONENT_EXPORT(WEBNN_SERVICE) WebNNTensorImpl
    : public WebNNObjectImpl<mojom::WebNNTensor,
                             blink::WebNNTensorToken,
                             mojo::AssociatedReceiver<mojom::WebNNTensor>> {
 public:
  // Similar to OnTaskRunnerDeleter, but waits for deletion to complete when
  // destruction occurs off the task runner's sequence. This belongs here
  // because it is specific to the representation pointers owned by
  // WebNNTensorImpl and reuses its wait helper.
  struct COMPONENT_EXPORT(WEBNN_SERVICE) OnTaskRunnerDeleterWithWait {
    explicit OnTaskRunnerDeleterWithWait(
        scoped_refptr<base::SequencedTaskRunner> task_runner);
    ~OnTaskRunnerDeleterWithWait();

    OnTaskRunnerDeleterWithWait(OnTaskRunnerDeleterWithWait&&);
    OnTaskRunnerDeleterWithWait& operator=(OnTaskRunnerDeleterWithWait&&);

    // For compatibility with std:: deleters.
    template <typename T>
    void operator()(const T* ptr) {
      if (!ptr) {
        return;
      }
      WebNNTensorImpl::RunOrPostTaskAndWaitOnSequenceInternal(
          task_runner_, base::BindOnce(&base::DeletePointer<const T>, ptr));
    }

   private:
    scoped_refptr<base::SequencedTaskRunner> task_runner_;
  };

  using RepresentationPtr = std::unique_ptr<gpu::WebNNTensorRepresentation,
                                            OnTaskRunnerDeleterWithWait>;

  WebNNTensorImpl(mojo::PendingAssociatedReceiver<mojom::WebNNTensor> receiver,
                  WebNNContextImpl& context,
                  mojom::TensorInfoPtr tensor_info);

  WebNNTensorImpl(mojo::PendingAssociatedReceiver<mojom::WebNNTensor> receiver,
                  WebNNContextImpl& context,
                  mojom::TensorInfoPtr tensor_info,
                  RepresentationPtr representation);

  WebNNTensorImpl(const WebNNTensorImpl&) = delete;
  WebNNTensorImpl& operator=(const WebNNTensorImpl&) = delete;

  OperandDataType data_type() const { return descriptor_.data_type(); }
  const std::vector<uint32_t>& shape() const { return descriptor_.shape(); }
  MLTensorUsage usage() const { return usage_; }

  size_t PackedByteLength() const { return descriptor_.PackedByteLength(); }
  size_t NumberOfElements() const { return descriptor_.NumberOfElements(); }

  bool IsValidWithDescriptor(const OperandDescriptor& descriptor) const;

  // This method will be called by `WriteTensor()` after the write info is
  // validated. A backend subclass should implement this method to write data
  // to a platform specific buffer.
  virtual void WriteTensorImpl(mojo_base::BigBuffer src_buffer) = 0;

  // Returns true if the tensor has been exported (e.g., to WebGPU)
  // and is not currently being accessed by WebNN.
  // Used to prevent concurrent access between WebNN and other consumers.
  bool is_exported() const {
    return representation_ && !representation_access_;
  }

  // Returns true if the tensor is an interop tensor.
  bool has_shared_image() const { return representation_ != nullptr; }

  // This method will be called by `ImportTensor()` or
  // `WebNNContext::CreateTensorFromMailbox()` for WebNN to begin access of the
  // platform-specific tensor as a shared image and then call
  // `ImportTensorImpl()` with that access. Returns true on success.
  bool ImportTensorInternal();

  // Destroys tensor shared image access and representation on their bound
  // sequences, and waits for teardown to complete before returning.
  // Called when the tensor is disconnected from the context or when the
  // context is destroyed.
  void DestroyAccessAndRepresentationAndWait();

 protected:
  ~WebNNTensorImpl() override;

  // This method will be called by `ReadTensor()` after the read info is
  // validated. A backend subclass should implement this method to read data
  // from a platform specific buffer.
  virtual void ReadTensorImpl(
      mojom::WebNNTensor::ReadTensorCallback callback) = 0;

  using ScopedAccessPtr =
      std::unique_ptr<gpu::WebNNTensorRepresentation::ScopedAccess,
                      OnTaskRunnerDeleter>;

  // Called by `ExportTensor()` after WebNN ends access of the
  // platform-specific tensor as a shared image.
  // Backend subclasses implement this to perform any necessary
  // device synchronization.
  virtual void ExportTensorImpl(ScopedAccessPtr access) = 0;

  // Called by `ImportTensorInternal()` after WebNN begins access of the
  // platform-specific tensor as a shared image.
  // Backend subclasses implement this to perform any necessary
  // device synchronization and store the access. Returns true on success.
  // On success, the subclass should assign `representation_access_` to
  // `access`. Must not post tasks itself; all main thread synchronization is
  // handled by `ImportTensorInternal()`.
  virtual bool ImportTensorImpl(ScopedAccessPtr access) = 0;

  // The `WebNNContextImpl` which owns and will outlive this object.
  const base::raw_ref<WebNNContextImpl> context_;

  // The shared image representation used to access the contents from shared
  // image. Only valid when usage has WebGPUInterop.
  RepresentationPtr representation_{nullptr,
                                    OnTaskRunnerDeleterWithWait(nullptr)};

  // Non-null only while WebNN holds exclusive access. Null if exported.
  ScopedAccessPtr representation_access_{nullptr, OnTaskRunnerDeleter(nullptr)};

 private:
  // Helper that runs a closure synchronously on a different sequence.
  // The caller blocks but the target sequence never blocks.
  // It is important the task does not post back to the current sequence, to
  // prevent deadlocks.
  static void RunOrPostTaskAndWaitOnSequenceInternal(
      scoped_refptr<base::SequencedTaskRunner> target,
      base::OnceClosure task);

  // mojom::WebNNTensor
  void ReadTensor(ReadTensorCallback callback) override;
  void WriteTensor(mojo_base::BigBuffer src_buffer) override;
  void ImportTensor(uint64_t flow_id, const gpu::SyncToken& fence) override;
  void ExportTensor(uint64_t flow_id, uint64_t release_count) override;
  void ExportTensorSync(uint64_t flow_id,
                        uint64_t release_count,
                        ExportTensorSyncCallback callback) override;

  // `OnDisconnect` is called from two places.
  //  - When the tensor is explicitly destroyed by the WebNN
  //  developer via the WebNN API.
  //  - When the tensor is dropped by the WebNN developer where
  //  the tensor gets implicitly destroyed upon garbage collection.
  void OnDisconnect() override;

  const OperandDescriptor descriptor_;
  const MLTensorUsage usage_;
};

}  // namespace webnn

#endif  // SERVICES_WEBNN_WEBNN_TENSOR_IMPL_H_