text: // Copyright (c) Microsoft Corporation. All rights reserved.
  // Licensed under the MIT License.
  
  #include "lib/Api/pch/pch.h"
  
  #include "LearningModel.h"
  
  #include "TelemetryEvent.h"
  #include "MapFeatureDescriptor.h"
  #include "SequenceFeatureDescriptor.h"
  #include "TensorFeatureDescriptor.h"
  
  #include "OnnxruntimeProvider.h"
  
  #include <robuffer.h>
  
  namespace WINMLP {
  
  // IBuffer implementation to avoid calling into WinTypes.dll to create wss::Buffer.
  // This will enable model creation on VTL1 without pulling in additional binaries on load.
  template <typename T>
  class STLVectorBackedBuffer : public winrt::implements<
                                    STLVectorBackedBuffer<T>,
                                    wss::IBuffer,
                                    ::Windows::Storage::Streams::IBufferByteAccess> {
   private:
    std::vector<T> data_;
    size_t length_ = 0;
  
   public:
    STLVectorBackedBuffer(size_t num_elements) : data_(num_elements) {}
  
    uint32_t Capacity() const try {
      // Return the size of the backing vector in bytes
      return static_cast<uint32_t>(data_.size() * sizeof(T));
    }
    WINML_CATCH_ALL
  
    uint32_t Length() const try {
      // Return the used buffer in bytes
      return static_cast<uint32_t>(length_);
    }
    WINML_CATCH_ALL
  
    void Length(uint32_t value) try {
      // Set the use buffer length in bytes
      WINML_THROW_HR_IF_TRUE_MSG(E_INVALIDARG, value > Capacity(), "Parameter 'value' cannot be greater than the buffer's capacity.");
      length_ = value;
    }
    WINML_CATCH_ALL
  
    STDMETHOD(Buffer)
    (_Outptr_ BYTE** value) {
      // Return the buffer
      RETURN_HR_IF_NULL(E_POINTER, value);
      *value = reinterpret_cast<BYTE*>(data_.data());
      return S_OK;
    }
  };
  
  
  LearningModel::LearningModel(
      const hstring& path,
      const winml::ILearningModelOperatorProvider op_provider) try : operator_provider_(op_provider) {
    _winmlt::TelemetryEvent loadModel_event(_winmlt::EventCategory::kModelLoad);
  
    WINML_THROW_IF_FAILED(CreateOnnxruntimeEngineFactory(engine_factory_.put()));
  
    wil::unique_handle file_handle{
  #if WINVER >= _WIN32_WINNT_WIN8
        CreateFile2(path.c_str(),
                    GENERIC_READ,
                    FILE_SHARE_READ,
                    OPEN_EXISTING,
                    NULL)};
  #else
        CreateFileW(path.c_str(),
                    GENERIC_READ,
                    FILE_SHARE_READ,
                    NULL,
                    OPEN_EXISTING,
                    FILE_ATTRIBUTE_READONLY,
                    NULL)};
  #endif
  
    WINML_THROW_HR_IF_TRUE_MSG(__HRESULT_FROM_WIN32(GetLastError()),
                               file_handle.get() == INVALID_HANDLE_VALUE,
                               "Model load failed!");
  
    auto file_mapping = wil::unique_handle(CreateFileMappingW(file_handle.get(),  // current file handle
                                                              NULL,               // default security
                                                              PAGE_READONLY,      // read/write permission
                                                              0,                  // size of mapping object, high
                                                              0,                  // size of mapping object, low
                                                              NULL));             // name of mapping object
  
    WINML_THROW_HR_IF_TRUE_MSG(__HRESULT_FROM_WIN32(GetLastError()),
                               file_mapping == nullptr,
                               "Model load failed!");
  
    auto buffer = MapViewOfFile(file_mapping.get(), // handle to mapping object
                                FILE_MAP_READ,      // read/write
                                0,                  // high-order 32 bits of file offset
                                0,                  // low-order 32 bits of file offset
                                0);                 // number of bytes to map. 0 means read whole file.
  
    WINML_THROW_HR_IF_TRUE_MSG(__HRESULT_FROM_WIN32(GetLastError()),
                               file_mapping == nullptr,
                               "Model load failed!");
    LARGE_INTEGER file_size;
    WINML_THROW_HR_IF_FALSE_MSG(__HRESULT_FROM_WIN32(GetLastError()),
                                GetFileSizeEx(file_handle.get(), &file_size),
                                "GetFileSizeEx");
    WINML_THROW_IF_FAILED(engine_factory_->CreateModel(buffer, static_cast<size_t>(file_size.QuadPart), model_.put()));
    WINML_THROW_HR_IF_TRUE_MSG(E_UNEXPECTED, UnmapViewOfFile(buffer) == 0, "Could not unmap model file.");
    WINML_THROW_IF_FAILED(model_->GetModelInfo(model_info_.put()));
  }
  WINML_CATCH_ALL
  
  LearningModel::LearningModel(
      _winml::IEngineFactory* engine_factory,
      _winml::IModel* model,
      const winml::ILearningModelOperatorProvider operator_provider) try :
        operator_provider_(operator_provider) {
    engine_factory_.copy_from(engine_factory);
    model_.copy_from(model);
    WINML_THROW_IF_FAILED(model_->GetModelInfo(model_info_.put()));
  }
  WINML_CATCH_ALL
  
  static HRESULT CreateModelFromStream(
      _winml::IEngineFactory* engine_factory,
      const wss::IRandomAccessStreamReference stream,
      _winml::IModel** model) {
    auto content = stream.OpenReadAsync().get();
  
    auto buffer = winrt::make<STLVectorBackedBuffer<BYTE>>(static_cast<size_t>(content.Size()));
    auto result = content.ReadAsync(
                             buffer,
                             buffer.Capacity(),
                             wss::InputStreamOptions::None).get();
  
    auto bytes = buffer.try_as<::Windows::Storage::Streams::IBufferByteAccess>();
    WINML_THROW_HR_IF_NULL_MSG(E_UNEXPECTED, bytes, "Model stream is invalid.");
  
    void* data;
    WINML_THROW_IF_FAILED_MSG(bytes->Buffer(reinterpret_cast<byte**>(&data)), "Failed to acquire buffer from model stream.");
  
    size_t len = static_cast<size_t>(content.Size());
    if (FAILED(engine_factory->CreateModel(data, len, model))) {
      WINML_THROW_HR(E_INVALIDARG);
    }
  
    return S_OK;
  }
  
  LearningModel::LearningModel(
      const wss::IRandomAccessStreamReference stream,
      const winml::ILearningModelOperatorProvider operator_provider) try : operator_provider_(operator_provider) {
    _winmlt::TelemetryEvent loadModel_event(_winmlt::EventCategory::kModelLoad);
  
    WINML_THROW_IF_FAILED(CreateOnnxruntimeEngineFactory(engine_factory_.put()));
    WINML_THROW_IF_FAILED(CreateModelFromStream(engine_factory_.get(), stream, model_.put()));
    WINML_THROW_IF_FAILED(model_->GetModelInfo(model_info_.put()));
  }
  WINML_CATCH_ALL
  
  static HRESULT CreateModelFromBuffer(
      _winml::IEngineFactory* engine_factory,
      const wss::IBuffer buffer,
      _winml::IModel** model) {
  		
    size_t len = buffer.Length();
    if (FAILED(engine_factory->CreateModel((void*)buffer.data(), len, model))) {
      WINML_THROW_HR(E_INVALIDARG);
    }
  
    return S_OK;
  }
  
  LearningModel::LearningModel(
      const wss::IBuffer buffer,
      const winml::ILearningModelOperatorProvider operator_provider) try : operator_provider_(operator_provider) {
    _winmlt::TelemetryEvent loadModel_event(_winmlt::EventCategory::kModelLoad);
  
    WINML_THROW_IF_FAILED(CreateOnnxruntimeEngineFactory(engine_factory_.put()));
    WINML_THROW_IF_FAILED(CreateModelFromBuffer(engine_factory_.get(), buffer, model_.put()));
    WINML_THROW_IF_FAILED(model_->GetModelInfo(model_info_.put()));
  }
  WINML_CATCH_ALL
  
  hstring
  LearningModel::Author() try {
    const char* out;
    size_t len;
    WINML_THROW_IF_FAILED(model_info_->GetAuthor(&out, &len));
    return _winml::Strings::HStringFromUTF8(out);
  }
  WINML_CATCH_ALL
  
  hstring
  LearningModel::Name() try {
    const char* out;
    size_t len;
    WINML_THROW_IF_FAILED(model_info_->GetName(&out, &len));
    return _winml::Strings::HStringFromUTF8(out);
  }
  WINML_CATCH_ALL
  
  hstring
  LearningModel::Domain() try {
    const char* out;
    size_t len;
    WINML_THROW_IF_FAILED(model_info_->GetDomain(&out, &len));
    return _winml::Strings::HStringFromUTF8(out);
  }
  WINML_CATCH_ALL
  
  hstring
  LearningModel::Description() try {
    const char* out;
    size_t len;
    WINML_THROW_IF_FAILED(model_info_->GetDescription(&out, &len));
    return _winml::Strings::HStringFromUTF8(out);
  }
  WINML_CATCH_ALL
  
  int64_t
  LearningModel::Version() try {
    int64_t version;
    WINML_THROW_IF_FAILED(model_info_->GetVersion(&version));
    return version;
  }
  WINML_CATCH_ALL
  
  wfc::IMapView<hstring, hstring>
  LearningModel::Metadata() try {
    ABI::Windows::Foundation::Collections::IMapView<HSTRING, HSTRING>* metadata = nullptr;
    wfc::IMapView<hstring, hstring> out;
    WINML_THROW_IF_FAILED(model_info_->GetModelMetadata(&metadata));
    winrt::attach_abi(out, metadata);
    return out;
  }
  WINML_CATCH_ALL
  
  IMLOperatorRegistry*
  LearningModel::GetOperatorRegistry() {
    if (operator_provider_ == nullptr) {
      return nullptr;
    }
  
    // Get the native winrt provider interface out of winrt operator provider.
    auto operator_provider_native =
        operator_provider_.as<ILearningModelOperatorProviderNative>();
  
    IMLOperatorRegistry* registry = nullptr;
    // Retrieve the "operator abi" registry.
    THROW_IF_FAILED(operator_provider_native->GetRegistry(&registry));
    return registry;
  }
  
  wfc::IVectorView<winml::ILearningModelFeatureDescriptor>
  LearningModel::InputFeatures() try {
    ABI::Windows::Foundation::Collections::IVectorView<winml::ILearningModelFeatureDescriptor>* features = nullptr;
    wfc::IVectorView<winml::ILearningModelFeatureDescriptor> out;
    WINML_THROW_IF_FAILED(model_info_->GetInputFeatures(&features));
    winrt::attach_abi(out, features);
    return out;
  }
  WINML_CATCH_ALL
  
  wfc::IVectorView<winml::ILearningModelFeatureDescriptor>
  LearningModel::OutputFeatures() try {
    ABI::Windows::Foundation::Collections::IVectorView<winml::ILearningModelFeatureDescriptor>* features = nullptr;
    wfc::IVectorView<winml::ILearningModelFeatureDescriptor> out;
    WINML_THROW_IF_FAILED(model_info_->GetOutputFeatures(&features));
    winrt::attach_abi(out, features);
    return out;
  }
  WINML_CATCH_ALL
  
  void LearningModel::SetName(const hstring& name) try {
    auto name_std_str = _winml::Strings::UTF8FromHString(name);
    auto name_c_str = name_std_str.c_str();
    WINML_THROW_IF_FAILED(model_->SetName(name_c_str));
  }
  WINML_CATCH_ALL
  
  void LearningModel::Close() try {
    // close the model
    model_ = nullptr;
  }
  WINML_CATCH_ALL
  
  bool LearningModel::IsDisposed() {
    return model_ == nullptr;
  }
  
  wf::IAsyncOperation<winml::LearningModel>
  LearningModel::LoadFromStorageFileAsync(
      ws::IStorageFile const modelFile) {
    return LoadFromStorageFileAsync(modelFile, nullptr);
  }
  
  wf::IAsyncOperation<winml::LearningModel>
  LearningModel::LoadFromStorageFileAsync(
      ws::IStorageFile const modelFile,
      winml::ILearningModelOperatorProvider const provider) {
    co_await resume_background();
    return make<LearningModel>(modelFile, provider);
  }
  
  wf::IAsyncOperation<winml::LearningModel>
  LearningModel::LoadFromStreamAsync(
      wss::IRandomAccessStreamReference const model_stream) {
    return LoadFromStreamAsync(model_stream, nullptr);
  }
  
  wf::IAsyncOperation<winml::LearningModel>
  LearningModel::LoadFromStreamAsync(
      wss::IRandomAccessStreamReference const model_stream,
      winml::ILearningModelOperatorProvider const provider) {
    co_await resume_background();
    return make<LearningModel>(model_stream, provider);
  }
  
  wf::IAsyncOperation<winml::LearningModel>
  LearningModel::LoadFromBufferAsync(
      wss::IBuffer const model_buffer) {
    return LoadFromBufferAsync(model_buffer, nullptr);
  }
  
  wf::IAsyncOperation<winml::LearningModel>
  LearningModel::LoadFromBufferAsync(
      wss::IBuffer const model_buffer,
      winml::ILearningModelOperatorProvider const provider) {
    co_await resume_background();
    return make<LearningModel>(model_buffer, provider);
  }
  
  winml::LearningModel
  LearningModel::LoadFromFilePath(
      hstring const& path) try {
    return LoadFromFilePath(path, nullptr);
  }
  WINML_CATCH_ALL
  
  winml::LearningModel
  LearningModel::LoadFromFilePath(
      hstring const& path,
      winml::ILearningModelOperatorProvider const provider) try {
    return make<LearningModel>(path, provider);
  }
  WINML_CATCH_ALL
  
  winml::LearningModel
  LearningModel::LoadFromStream(
      wss::IRandomAccessStreamReference const model_stream) try {
    return LoadFromStream(model_stream, nullptr);
  }
  WINML_CATCH_ALL
  
  winml::LearningModel
  LearningModel::LoadFromStream(
      wss::IRandomAccessStreamReference const model_stream,
      winml::ILearningModelOperatorProvider const provider) try {
    return make<LearningModel>(model_stream, provider);
  }
  WINML_CATCH_ALL
  
  winml::LearningModel
  LearningModel::LoadFromBuffer(
      wss::IBuffer const model_buffer) try {
    return LoadFromBuffer(model_buffer, nullptr);
  }
  WINML_CATCH_ALL
  
  winml::LearningModel
  LearningModel::LoadFromBuffer(
      wss::IBuffer const model_buffer,
      winml::ILearningModelOperatorProvider const provider) try {
    return make<LearningModel>(model_buffer, provider);
  }
  WINML_CATCH_ALL
  
  _winml::IModel*
  LearningModel::DetachModel() {
    com_ptr<_winml::IModel> detached_model;
    if (model_ != nullptr) {
      detached_model.attach(model_.detach());
  
      // Close the model since we now own the model proto
      Close();
    }
    return detached_model.detach();
  }
  
  _winml::IModel*
  LearningModel::CloneModel() {
    if (model_ == nullptr) {
      return nullptr;
    }
  
    com_ptr<_winml::IModel> model_copy;
    WINML_THROW_IF_FAILED(model_->CloneModel(model_copy.put()));
  
    return model_copy.detach();
  }
  
  _winml::IEngineFactory*
  LearningModel::GetEngineFactory() {
    return engine_factory_.get();
  }
  
  void LearningModel::SaveToFile(const hstring& file_name) {
    model_->SaveModel(file_name.c_str(), file_name.size());
  }
  
  void LearningModel::JoinModel(
      winml::LearningModel other,
      const std::unordered_map<std::string, std::string>& linkages,
      bool promote_unlinked_outputs,
      bool close_model_on_join,
      const winrt::hstring& join_node_prefix) {
    auto otherp = other.as<winmlp::LearningModel>();
    winrt::com_ptr<_winml::IModel> other_model;
    if (close_model_on_join) {
      other_model.attach(otherp->DetachModel());
    } else {
      other_model.attach(otherp->CloneModel());
    }
  
    std::vector<const char*> raw_outputs(linkages.size());
    std::vector<const char*> raw_inputs(linkages.size());
    std::transform(std::begin(linkages), std::end(linkages), std::begin(raw_outputs),
                   [](auto& pair) { return pair.first.c_str(); });
    std::transform(std::begin(linkages), std::end(linkages), std::begin(raw_inputs),
                   [](auto& pair) { return pair.second.c_str(); });
  
    auto prefix = winrt::to_string(join_node_prefix);
    WINML_THROW_IF_FAILED(model_->JoinModel(other_model.get(),
                                            raw_outputs.data(),
                                            raw_inputs.data(),
                                            linkages.size(),
                                            promote_unlinked_outputs,
                                            prefix.c_str()));
  
    model_info_ = nullptr;
    WINML_THROW_IF_FAILED(model_->GetModelInfo(model_info_.put()));
  }
  
  }  // namespace WINMLP
  
  namespace WINML::factory_implementation {
  // copied from cppwinrt magic to create abi wrappers.   Need to do it this way
  // since peeps underneath (like the constructor) will throw
  HRESULT
  __stdcall LearningModel::Load(
      const wchar_t* p_model_path,
      uint32_t model_path_size,
      IUnknown** pp_model_unk) {
    try {
      WINML_THROW_HR_IF_NULL_MSG(E_INVALIDARG, p_model_path, "Failed to create LearningModel. Ivalid argument p_model_path.");
      WINML_THROW_HR_IF_FALSE_MSG(E_INVALIDARG, model_path_size > 0, "Failed to create LearningModel. Ivalid argument model_path_size.");
      WINML_THROW_HR_IF_NULL_MSG(E_INVALIDARG, pp_model_unk, "Failed to create LearningModel. Ivalid argument pp_model_unk.");
  
      winrt::hstring path(p_model_path, model_path_size);
      auto model = make<winmlp::LearningModel>(path, nullptr);
      *pp_model_unk = model.as<IUnknown>().detach();
      return S_OK;
    }
    WINML_CATCH_ALL_COM
  }
  }  // namespace WINML::factory_implementation
  