// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#pragma once

#include "OnnxruntimeSessionBuilder.h"

namespace _winml {

class OnnxruntimeEngineFactory;

class OnnxruntimeDmlSessionBuilder : public Microsoft::WRL::RuntimeClass<
                                         Microsoft::WRL::RuntimeClassFlags<Microsoft::WRL::ClassicCom>,
                                         IOrtSessionBuilder> {
 public:
  HRESULT RuntimeClassInitialize(OnnxruntimeEngineFactory* engine_factory, ID3D12Device* device, ID3D12CommandQueue* queue, bool metacommands_enabled_);

  HRESULT STDMETHODCALLTYPE CreateSessionOptions(
      OrtSessionOptions** options) override;

  HRESULT STDMETHODCALLTYPE CreateSession(
      OrtSessionOptions* options,
      OrtThreadPool* inter_op_thread_pool,
      OrtThreadPool* intra_op_thread_pool,
      OrtSession** session) override;

  HRESULT STDMETHODCALLTYPE Initialize(
      OrtSession* session) override;

 private:
  Microsoft::WRL::ComPtr<OnnxruntimeEngineFactory> engine_factory_;
  winrt::com_ptr<ID3D12Device> device_;
  winrt::com_ptr<ID3D12CommandQueue> queue_;
  bool metacommands_enabled_ = true;
};

}  // namespace _winml