#ifndef THIRD_PARTY_UPB_WASM_ABI_HANDLE_H_
#define THIRD_PARTY_UPB_WASM_ABI_HANDLE_H_

#include <cstdint>
#include <variant>

#include "third_party/absl/base/attributes.h"
#include "third_party/absl/base/nullability.h"
#include "third_party/absl/base/optimization.h"
#include "third_party/absl/container/inlined_vector.h"
#include "third_party/absl/log/check.h"
#include "third_party/absl/status/status.h"
#include "third_party/absl/strings/str_cat.h"
#include "third_party/absl/strings/string_view.h"
#include "wasm_abi/handle_constants.h"
#include "wasm_abi/status_util.h"

// Must be last.
#include "wasm_abi/def.inc"

struct upb_Arena;
struct upb_Array;
struct upb_ExtensionRegistry;
struct upb_Map;
struct upb_Message;
struct upb_MiniTable;
struct upb_MiniTableEnum;
struct upb_MiniTableField;

namespace upb::wasm {

static inline constexpr uint32_t kNullHandleValue = 0xffffffff;

template <typename T>
struct HandleType {
  static constexpr uint32_t id() { return IdForHandle<T>(); }
  static constexpr absl::string_view debug_name() {
    return DebugStringForHandle<T>();
  }
  static constexpr uint32_t tag() { return (((id() & 0x0f) << 4) | 0x0f); }

  static constexpr uint32_t CreateTaggedHandle(uint32_t handle) {
    return (handle << 8) | tag();
  }
  static constexpr uint32_t null_handle() { return kNullHandleValue & tag(); }

  UPB_WASM_PRESERVE_ALL ABSL_ATTRIBUTE_NOINLINE static uint32_t
  HandleTaggedHandleError(MaybeErrorStatus* absl_nonnull status,
                          uint32_t tagged_handle, uint32_t max_handle) {
    uint32_t actual_tag = tagged_handle & 0xff;
    if ((actual_tag & 0x0f) != 0x0f) {
      upb::wasm::SetStatus(
          status, absl::StatusCode::kInvalidArgument,
          absl::StrCat("Handle ", absl::Hex(tagged_handle), " for type ",
                       debug_name(),
                       " misses the correct final four bits (0x0f)"));
      return kNullHandleValue;
    }
    if (actual_tag != tag()) {
      upb::wasm::SetStatus(
          status, absl::StatusCode::kInvalidArgument,
          absl::StrCat("Handle ", absl::Hex(tagged_handle), " has tag ",
                       absl::Hex(actual_tag), " but expected ",
                       absl::Hex(tag()), " for type ", debug_name(),
                       "; actual tag suggests ",
                       DebugStringForId(actual_tag >> 4)));
      return kNullHandleValue;
    }
    uint32_t handle = tagged_handle >> 8;
    if (handle <= 0) {
      upb::wasm::SetStatus(
          status, absl::StatusCode::kInvalidArgument,
          absl::StrCat("Invalid zero handle ", absl::Hex(handle), " (",
                       absl::Hex(tagged_handle), ") for type ", debug_name()));
      return kNullHandleValue;
    }
    if (handle > max_handle) {
      upb::wasm::SetStatus(status, absl::StatusCode::kInvalidArgument,
                           absl::StrCat("Handle ", absl::Hex(handle), " (",
                                        absl::Hex(tagged_handle),
                                        ") not valid for type ", debug_name()));
      return kNullHandleValue;
    }
    return handle;
  }

  UPB_WASM_PRESERVE_ALL ABSL_ATTRIBUTE_NOINLINE static void
  HandleTypeMismatchError(MaybeErrorStatus* absl_nonnull status,
                          uint32_t handle, uint32_t tagged_handle) {
    upb::wasm::SetStatus(
        status, absl::StatusCode::kInvalidArgument,
        absl::StrCat("Handle ", absl::Hex(handle), " (",
                     absl::Hex(tagged_handle), ") for type ", debug_name(),
                     " does not match value type"));
  }

  // Returns the handle without the tag, or an error if the handle is invalid or
  // exceeds the `max_handle`.
  //
  // This is one of the hot code paths, so we want to keep the happy path as
  // optimized as possible.
  ABSL_ATTRIBUTE_ALWAYS_INLINE static uint32_t RemoveTag(
      MaybeErrorStatus* absl_nonnull status, uint32_t tagged_handle,
      uint32_t max_handle) {
    uint32_t actual_tag = tagged_handle & 0xff;
    uint32_t handle = tagged_handle >> 8;

    bool is_valid_tag = actual_tag == tag();
    bool is_valid_handle = handle > 0 && handle <= max_handle;
    if (ABSL_PREDICT_FALSE(!is_valid_tag || !is_valid_handle)) {
      return HandleTaggedHandleError(status, tagged_handle, max_handle);
    }
    return handle;
  }
};

template <typename... T>
class HandleTable {
 public:
  using FieldType = std::variant<T...>;
  using StorageType = absl::InlinedVector<FieldType, 2>;

  template <typename U>
  uint32_t Insert(const U& value) {
    using Type = HandleType<U>;
    if (IsNullValue<U>(value)) {
      return Type::null_handle();
    }

    resources_.emplace_back(value);
    return Type::CreateTaggedHandle(resources_.size());
  }

  // Replaces the handle with a null value.
  template <typename U>
  void MarkNull(uint32_t tagged_handle) {
    using Type = HandleType<U>;
    if (ABSL_PREDICT_FALSE(tagged_handle == Type::null_handle())) {
      return;
    }
    absl::Status* status = nullptr;
    uint32_t handle =
        Type::RemoveTag(&status, tagged_handle, resources_.size());
    if (ABSL_PREDICT_FALSE(status != nullptr)) {
      return;
    }
    resources_[handle - 1] = NullValue<U>();
  }

  template <typename U>
  U ResolveHandle(MaybeErrorStatus* absl_nonnull status,
                  uint32_t tagged_handle) const {
    using Type = HandleType<U>;
    if (ABSL_PREDICT_FALSE(tagged_handle == Type::null_handle())) {
      return NullValue<U>();
    }

    uint32_t handle = Type::RemoveTag(status, tagged_handle, resources_.size());
    if (ABSL_PREDICT_FALSE(*status != nullptr)) {
      return NullValue<U>();
    }
    const FieldType& value = resources_[handle - 1];
    if (ABSL_PREDICT_FALSE(!std::holds_alternative<U>(value))) {
      Type::HandleTypeMismatchError(status, handle, tagged_handle);
      return NullValue<U>();
    }
    return std::get<U>(value);
  }

  StorageType::iterator begin() { return resources_.begin(); }
  StorageType::const_iterator begin() const { return resources_.begin(); }
  StorageType::iterator end() { return resources_.end(); }
  StorageType::const_iterator end() const { return resources_.end(); }

  StorageType::size_type size() const { return resources_.size(); }

  void Clear() { resources_.clear(); }

 private:
  StorageType resources_;
};

}  // namespace upb::wasm

#include "wasm_abi/undef.inc"
#endif  // THIRD_PARTY_UPB_WASM_ABI_HANDLE_H_
