#include "../../../devices/nvidia/nvidia_common.cuh"
#include "rope_nvidia.cuh"

#include "../../../devices/nvidia/nvidia_kernel_common.cuh"

#include "../cuda/kernel.cuh"

template <bool IsGPTJ, typename Tdata, typename Tindex, typename Tangle>
INFINIOP_CUDA_KERNEL ropeThreadPerItemKernel(
    Tdata *y_,
    const Tdata *x_,
    const Tindex *__restrict__ pos_ids,
    const Tangle *__restrict__ sin_table,
    const Tangle *__restrict__ cos_table,
    size_t table_dim,
    size_t pos_stride_batch, // Stride for batch dimension in pos_ids
    bool pos_has_batch_dim,  // Whether pos_ids has batch dimension
    bool has_batch_dim,      // Whether tensors have batch dimension
    ptrdiff_t y_stride_batch,
    ptrdiff_t y_stride_seqlen,
    ptrdiff_t y_stride_nhead,
    ptrdiff_t x_stride_batch,
    ptrdiff_t x_stride_seqlen,
    ptrdiff_t x_stride_nhead) {
    ropeThreadPerItemBlock<IsGPTJ>(
        y_, x_, pos_ids,
        sin_table, cos_table,
        table_dim,
        pos_stride_batch,
        pos_has_batch_dim,
        has_batch_dim,
        y_stride_batch, y_stride_seqlen, y_stride_nhead,
        x_stride_batch, x_stride_seqlen, x_stride_nhead);
}

template <bool IsGPTJ, typename Tindex, typename Tangle>
INFINIOP_CUDA_KERNEL customropeThreadPerItemKernel(
    cuda_bfloat16 *y_,
    const cuda_bfloat16 *x_,
    const Tindex *__restrict__ pos_ids,
    const Tangle *__restrict__ sin_table,
    const Tangle *__restrict__ cos_table,
    size_t table_dim,
    size_t pos_stride_batch, // Stride for batch dimension in pos_ids
    bool pos_has_batch_dim,  // Whether pos_ids has batch dimension
    bool has_batch_dim,      // Whether tensors have batch dimension
    ptrdiff_t y_stride_batch,
    ptrdiff_t y_stride_seqlen,
    ptrdiff_t y_stride_nhead,
    ptrdiff_t x_stride_batch,
    ptrdiff_t x_stride_seqlen,
    ptrdiff_t x_stride_nhead) {
    customropeThreadPerItemBlock<IsGPTJ>(
        y_, x_, pos_ids,
        sin_table, cos_table,
        table_dim,
        pos_stride_batch,
        pos_has_batch_dim,
        has_batch_dim,
        y_stride_batch, y_stride_seqlen, y_stride_nhead,
        x_stride_batch, x_stride_seqlen, x_stride_nhead);
}

namespace op::rope::nvidia {

struct Descriptor::Opaque {
    std::shared_ptr<device::nvidia::Handle::Internal> internal;
};

Descriptor::~Descriptor() {
    delete _opaque;
}

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::nvidia::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,
        new Opaque{reinterpret_cast<device::nvidia::Handle *>(handle)->internal()},
        handle->device,
        handle->device_id);

    return INFINI_STATUS_SUCCESS;
}

template <typename Tdata, typename Tindex>
infiniStatus_t calculateRoPE(const RoPEInfo &info,
                             int block_size,
                             Tdata *y,
                             const Tdata *x,
                             const Tindex *pos_ids,
                             const Tdata *sin_table,
                             const Tdata *cos_table,
                             cudaStream_t stream) {
    auto dimx = uint32_t(info.seqlen);
    auto dimy = uint32_t(info.nhead);

    // For 3D tensors, batch dimension is not in blockIdx
    auto dimz = info.has_batch_dim ? uint32_t(info.batch) : 1;

    int nthreads = std::max(int(info.table_dim), block_size);
    bool is_gpt_j = info.algo == infiniopRoPEAlgo_t::INFINIOP_ROPE_ALGO_GPT_J;

    // Calculate position ID stride for batch dimension
    size_t pos_stride_batch = info.pos_has_batch_dim ? info.seqlen : 0;

    dim3 grid_dim;
    if (info.has_batch_dim) {
        // 4D tensors: use 3D grid [seqlen, nhead, batch]
        grid_dim = dim3(dimx, dimy, dimz);
    } else {
        // 3D tensors: use 2D grid [seqlen, nhead], batch dimension is 1
        grid_dim = dim3(dimx, dimy, 1);
    }
    // printf("block_size = %d info.table_dim = %ld has_batch_dim: %d, is_gpt_j: %d pos_has_batch_dim: %d\n",
    //      block_size, info.table_dim, info.has_batch_dim, is_gpt_j, info.pos_has_batch_dim);
    // [batch, seqlen, nhead, dhead, table_len, table_dim, y_stride_batch, y_stride_seqlen, y_stride_nhead, x_stride_batch, x_stride_seqlen,x_stride_nhead]
    // printf("[%ld %ld %ld %ld %ld %ld %ld %ld %ld %ld %ld %ld]\n", info.batch,
    //      info.seqlen, info.nhead, info.dhead, info.table_len, info.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);
    if (is_gpt_j) {
        ropeThreadPerItemKernel<true><<<grid_dim, nthreads, 0, stream>>>(
            y, x, pos_ids, sin_table, cos_table, info.table_dim,
            pos_stride_batch,
            info.pos_has_batch_dim,
            info.has_batch_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);
    } else {
        if ((std::is_same<Tdata, cuda_bfloat16>::value) && (info.table_dim == 64) && (info.nhead < 16) && (info.seqlen < 505)) {
            auto bf16_y = reinterpret_cast<cuda_bfloat16*>(y); 
            auto bf16_x = reinterpret_cast<const cuda_bfloat16*>(x); 
            grid_dim = dim3(info.seqlen, info.batch, 1);
            dim3 block_dim = dim3(info.table_dim, info.nhead, 1); 
            customropeThreadPerItemKernel<false><<<grid_dim, block_dim, 0, stream>>>(
                bf16_y, bf16_x, pos_ids, sin_table, cos_table, info.table_dim,
                pos_stride_batch,
                info.pos_has_batch_dim,
                info.has_batch_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);
            // ropeThreadPerItemKernel<false><<<grid_dim, nthreads, 0, stream>>>(
            //     y, x, pos_ids, sin_table, cos_table, info.table_dim,
            //     pos_stride_batch,
            //     info.pos_has_batch_dim,
            //     info.has_batch_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);

        } else {
            ropeThreadPerItemKernel<false><<<grid_dim, nthreads, 0, stream>>>(
                y, x, pos_ids, sin_table, cos_table, info.table_dim,
                pos_stride_batch,
                info.pos_has_batch_dim,
                info.has_batch_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);
        }    
    }

    return INFINI_STATUS_SUCCESS;
}

#define CALCULATE_ROPE(TDATA, TINDEX)                      \
    calculateRoPE(_info,                                   \
                  _opaque->internal->maxThreadsPerBlock(), \
                  (TDATA *)y,                              \
                  (const TDATA *)x,                        \
                  (const TINDEX *)pos_ids,                 \
                  (const TDATA *)sin_table,                \
                  (const TDATA *)cos_table,                \
                  (cudaStream_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(cuda_bfloat16);
    case INFINI_DTYPE_F32:
        ROPE_TYPE(float);
    case INFINI_DTYPE_F64:
        ROPE_TYPE(double);
    default:
        return INFINI_STATUS_BAD_TENSOR_DTYPE;
    }

    return INFINI_STATUS_SUCCESS;
}

#undef ROPE_TYPE
#undef CALCULATE_ROPE

} // namespace op::rope::nvidia
