/*
 * Copyright 2023 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Static library loading in IREE.

#include "device/device.h"
#include "iree/hal/drivers/local_sync/sync_device.h"
#include "iree/hal/local/executable_plugin_manager.h"
#include "iree/hal/local/loaders/static_library_loader.h"
#include "model_util/model_api.h"

extern iree_hal_executable_import_provider_t importer_query(void);

// A function to create the HAL device from the different backend targets.
// The HAL device and loader are returned based on the implementation, and they
// must be released by the caller.
iree_status_t create_sample_device(iree_allocator_t host_allocator,
                                   iree_hal_device_t** out_device,
                                   iree_hal_executable_loader_t** loader) {
  iree_status_t status = iree_ok_status();

  // Set paramters for the device created in the next step.
  iree_hal_sync_device_params_t params;
  iree_hal_sync_device_params_initialize(&params);

  // Create plugin manager
  iree_hal_executable_plugin_manager_t* plugin_manager = NULL;

#ifdef BUILD_KELVIN
  // Enable plugin manager only in Kelvin for now
  if (iree_status_is_ok(status)) {
    status = iree_hal_executable_plugin_manager_create(
        /*capacity=*/1, host_allocator, &plugin_manager);
  }

  // Register import provider
  if (iree_status_is_ok(status)) {
    status = iree_hal_executable_plugin_manager_register_provider(
        plugin_manager, importer_query());
  }
#endif  // # ifdef BUILD_KELVIN

  // Load the statically embedded library
  const iree_hal_executable_library_query_fn_t libraries[] = {library_query()};

  if (iree_status_is_ok(status)) {
    status = iree_hal_static_library_loader_create(
        IREE_ARRAYSIZE(libraries), libraries,
        iree_hal_executable_plugin_manager_provider(plugin_manager),
        host_allocator, loader);
  }

  // Use the default host allocator for buffer allocations.
  iree_string_view_t identifier = iree_make_cstring_view("sync");
  iree_hal_allocator_t* device_allocator = NULL;
  if (iree_status_is_ok(status)) {
    status = iree_hal_allocator_create_heap(identifier, host_allocator,
                                            host_allocator, &device_allocator);
  }

  // Create the device and release the executor and loader afterwards.
  if (iree_status_is_ok(status)) {
    status = iree_hal_sync_device_create(
        identifier, &params, /*loader_count=*/1, loader, device_allocator,
        host_allocator, out_device);
  }

  iree_hal_allocator_release(device_allocator);
  return status;
}