#include "exception.h"
#include <cassert>
#include <cuda.h>
#include <tvm/ffi/extra/cuda/cubin_launcher.h>
#include <tvm/ffi/tvm_ffi.h>

#define CUDA_CHECK(code)                                                       \
  do {                                                                         \
    if ((code) != CUDA_SUCCESS) {                                              \
      throw triton_tvm_ffi::CUDAException(code);                               \
    }                                                                          \
  } while (false)

tvm::ffi::Map<tvm::ffi::String, int32_t> GetDeviceProperties(int device_id) {
  tvm::ffi::cuda_api::DeviceHandle device;
  CUDA_CHECK(cuDeviceGet(&device, device_id));
  int maxSharedMem = 0;
  int maxNumRegs = 0;
  int multiprocessorCount = 0;
  int warpSize = 0;
  int smClockRate = 0;
  int memClockRate = 0;
  int memBusWidth = 0;
  CUDA_CHECK(cuDeviceGetAttribute(
      &maxSharedMem, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN,
      device));
  CUDA_CHECK(cuDeviceGetAttribute(
      &maxNumRegs, CU_DEVICE_ATTRIBUTE_MAX_REGISTERS_PER_BLOCK, device));
  CUDA_CHECK(cuDeviceGetAttribute(
      &multiprocessorCount, CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, device));
  CUDA_CHECK(
      cuDeviceGetAttribute(&warpSize, CU_DEVICE_ATTRIBUTE_WARP_SIZE, device));
  CUDA_CHECK(cuDeviceGetAttribute(&smClockRate, CU_DEVICE_ATTRIBUTE_CLOCK_RATE,
                                  device));
  CUDA_CHECK(cuDeviceGetAttribute(
      &memClockRate, CU_DEVICE_ATTRIBUTE_MEMORY_CLOCK_RATE, device));
  CUDA_CHECK(cuDeviceGetAttribute(
      &memBusWidth, CU_DEVICE_ATTRIBUTE_GLOBAL_MEMORY_BUS_WIDTH, device));
  return {{"max_shared_mem", maxSharedMem},
          {"max_num_regs", maxNumRegs},
          {"multiprocessor_count", multiprocessorCount},
          {"warpSize", warpSize},
          {"sm_clock_rate", smClockRate},
          {"mem_clock_rate", memClockRate},
          {"mem_bus_width", memBusWidth}};
}

tvm::ffi::Tuple<uint64_t, uint64_t, int32_t, int32_t, int32_t>
LoadBinary(const tvm::ffi::String &name, const tvm::ffi::Bytes &data,
           int32_t shared, CUdevice device) {
  CUcontext pctx;
  CUfunction fun;
  CUmodule mod;
  int32_t nRegs = 0;
  int32_t nSpills = 0;
  int32_t nMaxThreads = 0;
  int32_t sharedOptin = 0;
  CUDA_CHECK(cuCtxGetCurrent(&pctx));
  if (!pctx) {
    CUDA_CHECK(cuDevicePrimaryCtxRetain(&pctx, device));
    CUDA_CHECK(cuCtxSetCurrent(pctx));
  }
  CUDA_CHECK(cuModuleLoadData(&mod, data.data()));
  CUDA_CHECK(cuModuleGetFunction(&fun, mod, name.data()));
  CUDA_CHECK(cuFuncGetAttribute(&nRegs, CU_FUNC_ATTRIBUTE_NUM_REGS, fun));
  CUDA_CHECK(
      cuFuncGetAttribute(&nSpills, CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, fun));
  CUDA_CHECK(cuFuncGetAttribute(&nMaxThreads,
                                CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK, fun));
  CUDA_CHECK(cuDeviceGetAttribute(
      &sharedOptin, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN,
      device));
  static constexpr int64_t kExpectedMaxDynamicSharedMemory = 49152;
  if (shared > kExpectedMaxDynamicSharedMemory &&
      sharedOptin > kExpectedMaxDynamicSharedMemory) {
    CUDA_CHECK(cuFuncSetCacheConfig(fun, CU_FUNC_CACHE_PREFER_SHARED));
    int32_t sharedTotal = 0, sharedStatic = 0;
    CUDA_CHECK(cuDeviceGetAttribute(
        &sharedTotal, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_MULTIPROCESSOR,
        device));
    CUDA_CHECK(cuFuncGetAttribute(&sharedStatic,
                                  CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, fun));
    CUDA_CHECK(
        cuFuncSetAttribute(fun, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
                           sharedOptin - sharedStatic));
  }
  return tvm::ffi::Tuple<uint64_t, uint64_t, int32_t, int32_t, int32_t>{
      mod, fun, nRegs, nSpills, nMaxThreads};
}

void Launch(tvm::ffi::PackedArgs args, tvm::ffi::Any *ret) {
  CUtensorMap x;
  int32_t gridX = args[0].cast<int32_t>();
  int32_t gridY = args[1].cast<int32_t>();
  int32_t gridZ = args[2].cast<int32_t>();
  CUstream stream = reinterpret_cast<CUstream>(args[3].cast<uint64_t>());
  CUfunction function = reinterpret_cast<CUfunction>(args[4].cast<uint64_t>());
  tvm::ffi::Tuple<int32_t, int32_t, int32_t> kernelMetadata =
      args[5].cast<tvm::ffi::Tuple<int32_t, int32_t, int32_t>>();
  int32_t numWarps = kernelMetadata.get<0>();
  int32_t numCtas = kernelMetadata.get<1>();
  int32_t sharedMemory = kernelMetadata.get<2>();
  tvm::ffi::ObjectRef launchMetadata = args[6].cast<tvm::ffi::ObjectRef>();
  tvm::ffi::ObjectRef launchEnterHook = args[7].cast<tvm::ffi::ObjectRef>();
  tvm::ffi::ObjectRef launchExitHook = args[8].cast<tvm::ffi::ObjectRef>();
  bool launchCooperativeGrid = args[9].cast<bool>();
  bool launchPdl = args[10].cast<bool>();
  tvm::ffi::ObjectRef globalScratchObject =
      args[11].cast<tvm::ffi::ObjectRef>();
  tvm::ffi::ObjectRef profileScratchObject =
      args[12].cast<tvm::ffi::ObjectRef>();
  tvm::ffi::PackedArgs kernelArgs = args.Slice(13);
  // TODO: call `launchEnterHook`
  // TODO: check `globalScratchObject`
  CUdeviceptr globalScratch = 0;
  // TODO: check `profileScratchObject`
  CUdeviceptr profileScratch = 0;
  if (gridX * gridY * gridZ > 0) {
    CUlaunchAttribute launchAttr[4];
    CUlaunchConfig config;
    config.gridDimX = gridX * numCtas;
    config.gridDimY = gridY;
    config.gridDimZ = gridZ;
    static constexpr int32_t kThreadsPerWarp = 32;
    config.blockDimX = kThreadsPerWarp * numWarps;
    config.blockDimY = 1;
    config.blockDimZ = 1;
    config.sharedMemBytes = sharedMemory;
    config.hStream = stream;
    config.attrs = launchAttr;
    int32_t numAttrs = 0;
    // TODO: check `launchPdf`
    // TODO: check `launchCooperativeGrid`
    if (numCtas != 1) {
      CUlaunchAttribute clusterAttr;
      clusterAttr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
      clusterAttr.value.clusterDim.x = numCtas;
      clusterAttr.value.clusterDim.y = 1;
      clusterAttr.value.clusterDim.z = 1;
      launchAttr[numAttrs++] = clusterAttr;
      CUlaunchAttribute clusterSchedulingAttr;
      clusterSchedulingAttr.id =
          CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE;
      clusterSchedulingAttr.value.clusterSchedulingPolicyPreference =
          CU_CLUSTER_SCHEDULING_POLICY_SPREAD;
      launchAttr[numAttrs++] = clusterSchedulingAttr;
    }
    config.numAttrs = numAttrs;
    if (numCtas == 16) {
      CUDA_CHECK(cuFuncSetAttribute(
          function, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1));
    }
    const int32_t kernelArgNum = kernelArgs.size();
    void **params =
        reinterpret_cast<void **>(alloca(sizeof(void *) * (kernelArgNum + 2)));
    for (size_t i = 0; i < kernelArgNum; ++i) {
      tvm::ffi::AnyView arg = kernelArgs[i];
      if (auto val = arg.try_cast<tvm::ffi::TensorView>()) {
        void **ptr = reinterpret_cast<void **>(alloca(sizeof(void *)));
        *ptr = val->data_ptr();
        params[i] = ptr;
      } else if (auto val = arg.try_cast<int32_t>()) {
        int32_t *ptr = reinterpret_cast<int32_t *>(alloca(sizeof(int32_t)));
        *ptr = *val;
        params[i] = ptr;
      } else if (auto val = arg.try_cast<float>()) {
        float *ptr = reinterpret_cast<float *>(alloca(sizeof(float)));
        *ptr = *val;
        params[i] = ptr;
      } else {
        assert(false && "unsupported kernel argument type");
      }
    }
    params[kernelArgNum] = &globalScratch;
    params[kernelArgNum + 1] = &profileScratch;
    CUDA_CHECK(cuLaunchKernelEx(&config, function, params, nullptr));
  }
  // TODO: call `launchExitHook`
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("triton_tvm_ffi.utils.get_device_properties", GetDeviceProperties)
      .def_packed("triton_tvm_ffi.utils.launch", Launch)
      .def("triton_tvm_ffi.utils.load_binary", LoadBinary);
}