#ifndef TRITON_TVM_FFI_TYPE_H_
#define TRITON_TVM_FFI_TYPE_H_

#include <cstdint>
#include <tvm/ffi/tvm_ffi.h>
#include <type_traits>

namespace triton_tvm_ffi {

// --------------- Definitions --------------- //

#define TYPE_TABLE_NATIVE(V)                                                   \
  V(I1, "i1", int8_t)                                                          \
  V(I8, "i8", int8_t)                                                          \
  V(I16, "i16", int16_t)                                                       \
  V(I32, "i32", int32_t)                                                       \
  V(I64, "i64", int64_t)                                                       \
  V(U1, "u1", uint8_t)                                                         \
  V(U8, "u8", uint8_t)                                                         \
  V(U16, "u16", uint16_t)                                                      \
  V(U32, "u32", uint32_t)                                                      \
  V(U64, "u64", uint64_t)                                                      \
  V(FP16, "fp16", double)                                                      \
  V(BF16, "bf16", double)                                                      \
  V(FP32, "f32", double)                                                       \
  V(FP64, "fp64", double)

#define TYPE_TABLE(V)                                                          \
  TYPE_TABLE_NATIVE(V)                                                         \
  V(PTR, "*?", void *)                                                         \
  V(CONSTEXPR, "constexpr", void)

enum class Type : int64_t {
#define DEFINE_ENUM(type, str, ctype) type,
  TYPE_TABLE(DEFINE_ENUM)
#undef DEFINE_ENUM
};

const char *TypeToString(Type type);
tvm::ffi::Optional<Type> StringToType(const tvm::ffi::String &name);
const char *TypeToCType(Type type);

template <Type T> struct type_to_ctype;
#define DEFINE_TYPE_TO_CTYPE(type, str, ctype)                                 \
  template <> struct type_to_ctype<Type::type> { using t = ctype; };
TYPE_TABLE(DEFINE_TYPE_TO_CTYPE)
#undef DEFINE_TYPE_TO_CTYPE
template <Type T> using type_to_ctype_t = typename type_to_ctype<T>::t;

// --------------- Implementations --------------- //

} // namespace triton_tvm_ffi

#endif