// Copyright 2023 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "iree/hal/api.h"

namespace iree::pjrt {

// Anonymous namespace containing helpers and wrappers for IREE API
// functions which can perform verbose logging when enabled. These all
// match an IREE api but will have the |iree_| prefix elided, so they are
// used as IreeApi::hal_allocator_allocate_buffer(...), which should be a
// drop-in for iree_hal_allocator_allocate_buffer(...).
namespace IreeApi {
namespace {

// Controls whether logging is printed to stderr. We may want to make this
// more configurable in the future.
const bool LOGGING_ENABLED = false;

IREE_PRINTF_ATTRIBUTE(2, 3)
void LogInvoke(const char* func, const char* fmt, ...) {
  if (LOGGING_ENABLED) {
    fprintf(stderr, ":: IREE INVOKE (%s): ", func);
    va_list args;
    va_start(args, fmt);
    vfprintf(stderr, fmt, args);
    va_end(args);
    fflush(stderr);
  }
}
iree_status_t HandleStatus(const char* func, iree_status_t status) {
  if (LOGGING_ENABLED) {
    if (!iree_status_is_ok(status)) {
      fprintf(stderr, " (");
      iree_status_fprint(stderr, status);
      fprintf(stderr, ")\n");
    } else {
      fprintf(stderr, " (OK)\n");
    }
  }
  return status;
}
std::string SemaphoreListToString(const iree_hal_semaphore_list_t sl) {
  std::string result;
  char fmtBuffer[64];
  for (iree_host_size_t i = 0; i < sl.count; ++i) {
    snprintf(fmtBuffer, sizeof(fmtBuffer), "%p:%" PRIu64, sl.semaphores[i],
             sl.payload_values[i]);
    if (i > 0) {
      result.append(", ");
    }
    result.append(fmtBuffer);
  }
  return result;
}
std::string FenceToString(iree_hal_fence_t* fence) {
  return SemaphoreListToString(iree_hal_fence_semaphore_list(fence));
}

iree_status_t hal_allocator_allocate_buffer(
    iree_hal_allocator_t* IREE_RESTRICT allocator,
    iree_hal_buffer_params_t params, iree_device_size_t allocation_size,
    iree_hal_buffer_t** out_buffer) {
  auto status = iree_hal_allocator_allocate_buffer(allocator, params,
                                                   allocation_size, out_buffer);
  if (LOGGING_ENABLED) {
    LogInvoke(__func__, "allocator=%p, size=%zu, buffer=%p", allocator,
              (size_t)allocation_size, *out_buffer);
  }
  return HandleStatus(__func__, status);
}

iree_status_t hal_allocator_import_buffer(
    iree_hal_allocator_t* IREE_RESTRICT allocator,
    iree_hal_buffer_params_t params,
    iree_hal_external_buffer_t* IREE_RESTRICT external_buffer,
    iree_hal_buffer_release_callback_t release_callback,
    iree_hal_buffer_t** out_buffer) {
  if (LOGGING_ENABLED) {
    LogInvoke(__func__, "external_buffer=%p", external_buffer);
  }
  return HandleStatus(__func__, iree_hal_allocator_import_buffer(
                                    allocator, params, external_buffer,
                                    release_callback, out_buffer));
}

iree_status_t hal_device_queue_alloca(
    iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
    const iree_hal_semaphore_list_t wait_semaphore_list,
    const iree_hal_semaphore_list_t signal_semaphore_list,
    iree_hal_pool_t* pool, iree_hal_buffer_params_t params,
    iree_device_size_t allocation_size, iree_hal_alloca_flags_t flags,
    iree_hal_buffer_t** IREE_RESTRICT out_buffer) {
  if (LOGGING_ENABLED) {
    LogInvoke(__func__, "device=%p, size=%zd, wait={%s}, signal={%s}", device,
              (size_t)allocation_size,
              SemaphoreListToString(wait_semaphore_list).c_str(),
              SemaphoreListToString(signal_semaphore_list).c_str());
  }
  return HandleStatus(__func__, iree_hal_device_queue_alloca(
                                    device, queue_affinity, wait_semaphore_list,
                                    signal_semaphore_list, pool, params,
                                    allocation_size, flags, out_buffer));
}

iree_status_t hal_device_queue_dealloca(
    iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
    const iree_hal_semaphore_list_t wait_semaphore_list,
    const iree_hal_semaphore_list_t signal_semaphore_list,
    iree_hal_buffer_t* buffer, iree_hal_dealloca_flags_t flags) {
  if (LOGGING_ENABLED) {
    LogInvoke(__func__, "device=%p, buffer=%p, wait={%s}, signal={%s}", device,
              buffer, SemaphoreListToString(wait_semaphore_list).c_str(),
              SemaphoreListToString(signal_semaphore_list).c_str());
  }
  return HandleStatus(__func__, iree_hal_device_queue_dealloca(
                                    device, queue_affinity, wait_semaphore_list,
                                    signal_semaphore_list, buffer, flags));
}

iree_status_t hal_device_queue_barrier(
    iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
    const iree_hal_semaphore_list_t wait_semaphore_list,
    const iree_hal_semaphore_list_t signal_semaphore_list) {
  if (LOGGING_ENABLED) {
    LogInvoke(__func__, "device=%p, wait={%s}, signal={%s}", device,
              SemaphoreListToString(wait_semaphore_list).c_str(),
              SemaphoreListToString(signal_semaphore_list).c_str());
  }
  return HandleStatus(__func__,
                      iree_hal_device_queue_barrier(
                          device, queue_affinity, wait_semaphore_list,
                          signal_semaphore_list, IREE_HAL_EXECUTE_FLAG_NONE));
}

iree_status_t hal_device_queue_execute(
    iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
    const iree_hal_semaphore_list_t wait_semaphore_list,
    const iree_hal_semaphore_list_t signal_semaphore_list,
    iree_hal_command_buffer_t* command_buffer) {
  if (LOGGING_ENABLED) {
    LogInvoke(__func__, "device=%p, wait={%s}, signal={%s}", device,
              SemaphoreListToString(wait_semaphore_list).c_str(),
              SemaphoreListToString(signal_semaphore_list).c_str());
  }
  return HandleStatus(__func__, iree_hal_device_queue_execute(
                                    device, queue_affinity, wait_semaphore_list,
                                    signal_semaphore_list, command_buffer,
                                    iree_hal_buffer_binding_table_empty(),
                                    IREE_HAL_EXECUTE_FLAG_NONE));
}

iree_status_t hal_fence_create(iree_host_size_t capacity,
                               iree_allocator_t host_allocator,
                               iree_hal_fence_t** out_fence) {
  auto status = iree_hal_fence_create(capacity, host_allocator, out_fence);
  if (LOGGING_ENABLED) {
    LogInvoke(__func__, "capacity=%zu, fence=%p", (size_t)capacity, *out_fence);
  }
  return HandleStatus(__func__, status);
}

iree_status_t hal_fence_create_at(iree_hal_semaphore_t* semaphore,
                                  uint64_t value,
                                  iree_allocator_t host_allocator,
                                  iree_hal_fence_t** out_fence) {
  auto status =
      iree_hal_fence_create_at(semaphore, value, host_allocator, out_fence);
  if (LOGGING_ENABLED) {
    LogInvoke(__func__, "semaphore=%p, value=%" PRIu64 ", fence=%p", semaphore,
              value, *out_fence);
  }
  return HandleStatus(__func__, status);
}

iree_status_t hal_fence_extend(iree_hal_fence_t* into_fence,
                               iree_hal_fence_t* from_fence) {
  if (LOGGING_ENABLED) {
    LogInvoke(__func__, "into_fence=%p, from_fence=%p", into_fence, from_fence);
  }
  return HandleStatus(__func__, iree_hal_fence_extend(into_fence, from_fence));
}

iree_status_t hal_fence_insert(iree_hal_fence_t* fence,
                               iree_hal_semaphore_t* semaphore,
                               uint64_t value) {
  if (LOGGING_ENABLED) {
    LogInvoke(__func__, "fence=%p, semaphore=%p, value=%" PRIu64, fence,
              semaphore, value);
  }
  return HandleStatus(__func__, iree_hal_fence_insert(fence, semaphore, value));
}

}  // namespace
}  // namespace IreeApi

}  // namespace iree::pjrt
