// 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_WEBNN_CONTEXT_PROVIDER_IN_RENDERER_H_
#define SERVICES_WEBNN_WEBNN_CONTEXT_PROVIDER_IN_RENDERER_H_

#include "base/component_export.h"
#include "base/files/file.h"
#include "base/memory/scoped_refptr.h"
#include "base/memory/weak_ptr.h"
#include "base/task/single_thread_task_runner.h"
#include "mojo/public/cpp/bindings/pending_remote.h"
#include "mojo/public/cpp/bindings/shared_remote.h"
#include "services/webnn/public/mojom/webnn_context_provider.mojom.h"
#include "services/webnn/webnn_context_provider_impl.h"

namespace webnn {

// A lightweight WebNNContextProvider implementation that runs without GPU
// dependencies (e.g., in the renderer process).
class COMPONENT_EXPORT(WEBNN_SERVICE) WebNNContextProviderInRenderer
    : public mojom::WebNNContextProvider {
 public:
  WebNNContextProviderInRenderer(
      mojo::PendingRemote<mojom::WebNNWeightsFileCreator>
          weights_file_creator_remote,
      scoped_refptr<base::SingleThreadTaskRunner> main_task_runner);
  ~WebNNContextProviderInRenderer() override;

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

  // mojom::WebNNContextProvider:
  void CreateWebNNContext(mojom::CreateContextOptionsPtr options,
                          CreateWebNNContextCallback callback) override;

  // Opens a writable tempfile in the browser and returns a per-file session
  // pipe. The renderer writes weights incrementally (calling
  // `RequestCapacityChange` before each extending write) and seals the file
  // with `Finalize` when done.
  void OpenWeightsFile(
      base::OnceCallback<void(base::File,
                              mojo::PendingRemote<mojom::WeightsFileSession>)>
          callback);

  base::WeakPtr<WebNNContextProviderInRenderer> GetWeakPtr() {
    return weak_factory_.GetWeakPtr();
  }

  void RemoveWebNNContextImpl(const blink::WebNNContextToken& handle);

 private:
  void OnCreateWebNNContextImpl(
      CreateWebNNContextCallback callback,
      mojo::PendingRemote<mojom::WebNNContext> remote,
      WebNNContextImpl::WebNNContextImplPtr context_impl);

  // SharedRemote for creating weights files in the browser process.
  // Uses SharedRemote so it can be called from any sequence.
  mojo::SharedRemote<mojom::WebNNWeightsFileCreator>
      shared_weights_file_creator_;

  // Task runner for the main thread (where the Mojo pipe lives).
  scoped_refptr<base::SingleThreadTaskRunner> main_task_runner_;

  // Contexts created by this provider. Cleaned up when the provider is
  // destroyed (when the mojo pipe closes).
  WebNNContextProviderImpl::WebNNContextImplSet context_impls_;

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

}  // namespace webnn

#endif  // SERVICES_WEBNN_WEBNN_CONTEXT_PROVIDER_IN_RENDERER_H_