#include "mediapipe/gpu/gpu_buffer.h"

#include <memory>
#include <utility>

#include "absl/functional/bind_front.h"
#include "absl/log/absl_check.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "mediapipe/framework/port/logging.h"

#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
#include "mediapipe/objc/util.h"
#endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER

namespace mediapipe {

namespace {

struct StorageTypeFormatter {
  void operator()(std::string* out,
                  const std::shared_ptr<internal::GpuBufferStorage>& s) const {
    absl::StrAppend(out, s->storage_type().name());
  }
};

}  // namespace

std::string GpuBuffer::DebugString() const {
  return holder_ ? absl::StrCat("GpuBuffer[", width(), "x", height(), " ",
                                format(), " as ", holder_->DebugString(), "]")
                 : "GpuBuffer[invalid]";
}

std::string GpuBuffer::StorageHolder::DebugString() const {
  absl::MutexLock lock(&mutex_);
  return absl::StrJoin(storages_, ", ", StorageTypeFormatter());
}

internal::GpuBufferStorage* GpuBuffer::StorageHolder::GetStorageForView(
    TypeId view_provider_type, bool for_writing) const {
  std::shared_ptr<internal::GpuBufferStorage> chosen_storage;
  std::function<std::shared_ptr<internal::GpuBufferStorage>()> conversion;

  {
    absl::MutexLock lock(&mutex_);
    // First see if any current storage supports the view.
    for (const auto& s : storages_) {
      if (s->can_down_cast_to(view_provider_type)) {
        chosen_storage = s;
        break;
      }
    }

    // Then try to convert existing storages to one that does.
    // TODO: choose best conversion.
    if (!chosen_storage) {
      for (const auto& s : storages_) {
        if (auto converter = internal::GpuBufferStorageRegistry::Get()
                                 .StorageConverterForViewProvider(
                                     view_provider_type, s->storage_type())) {
          conversion = absl::bind_front(converter, s);
          break;
        }
      }
    }
  }

  // Avoid invoking a converter or factory while holding the mutex.
  // Two reasons:
  // 1. Readers that don't need a conversion will not be blocked.
  // 2. We use mutexes to make sure GL contexts are not used simultaneously on
  //    different threads, and we also rely on Mutex's deadlock detection
  //    heuristic, which enforces a consistent mutex acquisition order.
  //    This function is likely to be called within a GL context, and the
  //    conversion function may in turn use a GL context, and this may cause a
  //    false positive in the deadlock detector.
  //    TODO: we could use Mutex::ForgetDeadlockInfo instead.
  if (conversion) {
    auto new_storage = conversion();
    absl::MutexLock lock(&mutex_);
    // Another reader might have already completed and inserted the same
    // conversion. TODO: prevent this?
    for (const auto& s : storages_) {
      if (s->can_down_cast_to(view_provider_type)) {
        chosen_storage = s;
        break;
      }
    }
    if (!chosen_storage) {
      storages_.push_back(std::move(new_storage));
      chosen_storage = storages_.back();
    }
  }

  if (for_writing) {
    // This will temporarily hold storages to be released, and do so while the
    // lock is not held (see above).
    decltype(storages_) old_storages;
    using std::swap;
    if (chosen_storage) {
      // Discard all other storages.
      absl::MutexLock lock(&mutex_);
      swap(old_storages, storages_);
      storages_ = {chosen_storage};
    } else {
      // Allocate a new storage supporting the requested view.
      if (auto factory =
              internal::GpuBufferStorageRegistry::Get()
                  .StorageFactoryForViewProvider(view_provider_type)) {
        if (auto new_storage = factory(width_, height_, format_)) {
          absl::MutexLock lock(&mutex_);
          swap(old_storages, storages_);
          storages_ = {std::move(new_storage)};
          chosen_storage = storages_.back();
        }
      }
    }
  }

  // It is ok to return a non-owning storage pointer here because this object
  // ensures the storage's lifetime. Overwriting a GpuBuffer while readers are
  // active would violate this, but it's not allowed in MediaPipe.
  return chosen_storage ? chosen_storage.get() : nullptr;
}

internal::GpuBufferStorage& GpuBuffer::GetStorageForViewOrDie(
    TypeId view_provider_type, bool for_writing) const {
  auto* chosen_storage =
      GpuBuffer::GetStorageForView(view_provider_type, for_writing);
  ABSL_CHECK(chosen_storage)
      << "no view provider found for requested view "
      << view_provider_type.name() << "; storages available: "
      << (holder_ ? holder_->DebugString() : "invalid");
  ABSL_DCHECK(chosen_storage->can_down_cast_to(view_provider_type));
  return *chosen_storage;
}

#if !MEDIAPIPE_DISABLE_GPU
#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
CVPixelBufferRef GetCVPixelBufferRef(const GpuBuffer& buffer) {
  if (buffer.GetStorageForView(
          kTypeId<internal::ViewProvider<CVPixelBufferRef>>,
          /*for_writing=*/false) != nullptr) {
    return *buffer.GetReadView<CVPixelBufferRef>();
  }
  return nullptr;
}
#endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
#endif  // !MEDIAPIPE_DISABLE_GPU

}  // namespace mediapipe