#include "../../../devices/bang/common_bang.h"
#include "rope_bang.h"
#include "rope_bang_kernel.mlu"

namespace op::rope::bang {

Descriptor::~Descriptor() = default;

infiniStatus_t Descriptor::create(
    infiniopHandle_t handle_,
    Descriptor **desc_ptr,
    infiniopTensorDescriptor_t y_desc,
    infiniopTensorDescriptor_t x_desc,
    infiniopTensorDescriptor_t pos_desc,
    infiniopTensorDescriptor_t sin_desc,
    infiniopTensorDescriptor_t cos_desc,
    infiniopRoPEAlgo_t algo) {

    auto handle = reinterpret_cast<device::bang::Handle *>(handle_);

    auto info = RoPEInfo::createRoPEInfo(y_desc, x_desc, pos_desc, sin_desc, cos_desc, algo);
    CHECK_RESULT(info);

    // Create descriptor
    *desc_ptr = new Descriptor(
        info.take(),
        0,
        nullptr,
        handle->device,
        handle->device_id);

    return INFINI_STATUS_SUCCESS;
}

template <typename Tdata, typename Tindex>
infiniStatus_t calculateRoPE(const RoPEInfo &info,
                             Tdata *y,
                             const Tdata *x,
                             const Tindex *pos_ids,
                             const Tdata *sin_table,
                             const Tdata *cos_table,
                             cnrtQueue_t queue) {
    auto batch_size = uint32_t(info.batch);
    auto seqlen = uint32_t(info.seqlen);
    auto nhead = uint32_t(info.nhead);
    auto table_dim = uint32_t(info.table_dim);

    cnrtDim3_t k_dim;
    cnrtFunctionType_t k_type;

    // Configure kernel launch parameters
    k_dim.x = 4;
    k_dim.y = 1;
    k_dim.z = 1;
    k_type = cnrtFuncTypeUnion1;

    // Launch kernel with batch dimension
    ropeKernel<<<k_dim, k_type, queue>>>(
        y, x, pos_ids, sin_table, cos_table,
        batch_size, seqlen, nhead, table_dim,
        info.y_stride_batch, info.y_stride_seqlen, info.y_stride_nhead,
        info.x_stride_batch, info.x_stride_seqlen, info.x_stride_nhead,
        info.algo);

    cnrtQueueSync(queue);

    return INFINI_STATUS_SUCCESS;
}

#define CALCULATE_ROPE(TDATA, TINDEX)       \
    calculateRoPE(_info,                    \
                  (TDATA *)y,               \
                  (const TDATA *)x,         \
                  (const TINDEX *)pos_ids,  \
                  (const TDATA *)sin_table, \
                  (const TDATA *)cos_table, \
                  (cnrtQueue_t)stream)

#define ROPE_TYPE(TDATA)                        \
    switch (_info.pos_type) {                   \
    case INFINI_DTYPE_U8:                       \
        return CALCULATE_ROPE(TDATA, uint8_t);  \
    case INFINI_DTYPE_U16:                      \
        return CALCULATE_ROPE(TDATA, uint16_t); \
    case INFINI_DTYPE_U32:                      \
        return CALCULATE_ROPE(TDATA, uint32_t); \
    case INFINI_DTYPE_U64:                      \
        return CALCULATE_ROPE(TDATA, uint64_t); \
    case INFINI_DTYPE_I8:                       \
        return CALCULATE_ROPE(TDATA, int8_t);   \
    case INFINI_DTYPE_I16:                      \
        return CALCULATE_ROPE(TDATA, int16_t);  \
    case INFINI_DTYPE_I32:                      \
        return CALCULATE_ROPE(TDATA, int32_t);  \
    case INFINI_DTYPE_I64:                      \
        return CALCULATE_ROPE(TDATA, int64_t);  \
    default:                                    \
        return INFINI_STATUS_BAD_TENSOR_DTYPE;  \
    }

infiniStatus_t Descriptor::calculate(
    void *workspace,
    size_t workspace_size,
    void *y,
    const void *x,
    const void *pos_ids,
    const void *sin_table,
    const void *cos_table,
    void *stream) const {

    switch (_info.data_type) {
    case INFINI_DTYPE_F16:
        ROPE_TYPE(half);
    case INFINI_DTYPE_BF16:
        ROPE_TYPE(bfloat16_t);
    case INFINI_DTYPE_F32:
        ROPE_TYPE(float);
    default:
        return INFINI_STATUS_BAD_TENSOR_DTYPE;
    }

    return INFINI_STATUS_SUCCESS;
}

#undef ROPE_TYPE
#undef CALCULATE_ROPE

} // namespace op::rope::bang
