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

#include "base/containers/queue.h"
#include "components/optimization_guide/core/model_execution/on_device_capability.h"
#include "components/optimization_guide/proto/on_device_model_execution_config.pb.h"
#include "services/on_device_model/public/mojom/on_device_model.mojom.h"

// Execution set for Optimization Guide sessions. It handles queueing requests
// for `ExecuteModel()` since multiple executions are not supported currently.
// Not thread-safe.
//
// TODO(crbug.com/403352722): When Optimization Guide supports broker interface
// via mojo, the logic in this class can be moved to the implementation of
// each feature.
class AIOnDeviceSession {
 public:
  explicit AIOnDeviceSession(
      std::unique_ptr<optimization_guide::OnDeviceSession> session);

  ~AIOnDeviceSession();

  // Not copyable or movable.
  AIOnDeviceSession(const AIOnDeviceSession&) = delete;
  AIOnDeviceSession& operator=(const AIOnDeviceSession&) = delete;

  // Queues the request for `OnDeviceSession::ExecuteModel()`.
  void ExecuteModelOrQueue(
      optimization_guide::MultimodalMessage request,
      optimization_guide::OptimizationGuideModelExecutionResultStreamingCallback
          callback,
      on_device_model::mojom::ResponseConstraintPtr constraint = nullptr);

  optimization_guide::OnDeviceSession* session() { return session_.get(); }

 private:
  // Takes the next pending request, if there is no execution in flight.
  void MaybeRunNextExecutionRequest();

  // Callback function invoked by `OnDeviceSession::ExecuteModel()`.
  void ModelExecutionCallback(
      optimization_guide::OptimizationGuideModelExecutionResultStreamingCallback
          final_callback,
      optimization_guide::OptimizationGuideModelStreamingExecutionResult
          result);

  // The underlying session provided by Optimization Guide.
  std::unique_ptr<optimization_guide::OnDeviceSession> session_;

  // Queue holding execution requests.
  struct ExecutionRequest {
    ExecutionRequest(
        optimization_guide::MultimodalMessage message,
        optimization_guide::
            OptimizationGuideModelExecutionResultStreamingCallback callback,
        on_device_model::mojom::ResponseConstraintPtr constraint);
    ExecutionRequest(ExecutionRequest&&);
    ~ExecutionRequest();

    optimization_guide::MultimodalMessage message;
    optimization_guide::OptimizationGuideModelExecutionResultStreamingCallback
        callback;
    on_device_model::mojom::ResponseConstraintPtr constraint;
  };
  base::queue<ExecutionRequest> requests_;

  bool is_execution_in_progress_ = false;

  base::WeakPtrFactory<AIOnDeviceSession> weak_ptr_factory_{this};
};

#endif  // CHROME_BROWSER_AI_AI_ON_DEVICE_SESSION_H_