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

#ifndef THIRD_PARTY_BLINK_RENDERER_MODULES_ML_WEBNN_ML_GRAPH_H_
#define THIRD_PARTY_BLINK_RENDERER_MODULES_ML_WEBNN_ML_GRAPH_H_

#include "base/types/pass_key.h"
#include "services/webnn/public/cpp/operand_descriptor.h"
#include "third_party/blink/renderer/bindings/core/v8/script_promise_resolver.h"
#include "third_party/blink/renderer/bindings/modules/v8/v8_ml_device_type.h"
#include "third_party/blink/renderer/bindings/modules/v8/v8_ml_operand_descriptor.h"
#include "third_party/blink/renderer/modules/modules_export.h"
#include "third_party/blink/renderer/platform/bindings/script_wrappable.h"
#include "third_party/blink/renderer/platform/heap/collection_support/heap_vector.h"
#include "third_party/blink/renderer/platform/heap/member.h"
#include "third_party/blink/renderer/platform/heap/visitor.h"

namespace blink {

class MLTensor;
class MLContext;
class MLGraphBuilder;
class ExecutionContext;

typedef HeapVector<std::pair<String, Member<MLTensor>>> MLNamedTensors;

// Represents a handle to a compiled, platform-specific computational graph.
class MODULES_EXPORT MLGraph : public ScriptWrappable {
  DEFINE_WRAPPERTYPEINFO();
  USING_PRE_FINALIZER(MLGraph, Dispose);

 public:
  using NamedOperandDescriptors =
      HashMap<String, std::optional<webnn::OperandDescriptor>>;

  // Instances should only be constructed via `MLGraphBuilder.build()`.
  // This method is public as required by the `MakeGarbageCollected` helper.
  //
  // `input_constraints` and `output_constraints` describe the constraints on
  // the inputs and outputs which may be used to execute the respective graph.
  MLGraph(ExecutionContext* execution_context,
          MLContext* context,
          blink::WebNNGraphToken graph_token,
          NamedOperandDescriptors input_constraints,
          NamedOperandDescriptors output_constraints,
          Vector<V8MLDeviceType> devices,
          base::PassKey<MLGraphBuilder> pass_key);

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

  ~MLGraph() override;

  void Trace(Visitor* visitor) const override;

  // ml_graph.idl
  void destroy();
  Vector<V8MLDeviceType> devices() const;

  const NamedOperandDescriptors& GetInputConstraints() const;
  const NamedOperandDescriptors& GetOutputConstraints() const;

  const blink::WebNNGraphToken& graph_token() const { return graph_token_; }

  bool IsDestroyed() const;

  const MLContext* Context() const;

 private:
  void Dispose();

  // Describes the constraints on the inputs or outputs to this graph.
  // Note that `HashMap` values must be nullable, but
  // `webnn::OperandDescriptor` lacks a default constructor, so an optional is
  // used. Do not add std::nullopt values to these maps.
  const NamedOperandDescriptors input_constraints_;
  const NamedOperandDescriptors output_constraints_;

  Member<MLContext> ml_context_;

  // Token identifying this graph for Dispatch calls on the context.
  const blink::WebNNGraphToken graph_token_;

  // Whether this graph has been destroyed.
  bool is_destroyed_ = false;

  // Devices that will be used when dispatching the graph.
  Vector<V8MLDeviceType> devices_;
};

}  // namespace blink

#endif  // THIRD_PARTY_BLINK_RENDERER_MODULES_ML_WEBNN_ML_GRAPH_H_