#include "../../../devices/metax/metax_common.h"
#include "kv_caching_metax.h"

#include "../../../devices/metax/metax_kernel_common.h"
#include <cub/block/block_reduce.cuh>

#include "../../../reduce/cuda/reduce.cuh"

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

template <typename Tdata, typename Tidx>
INFINIOP_METAX_KERNEL kvCaching(
    Tdata *k_cache,
    Tdata *v_cache,
    const Tdata *k,
    const Tdata *v,
    const Tidx *past_kv_lengths,
    int batch_size,
    int num_kv_heads,
    int max_seq_len,
    int seq_len,
    int hidden_dim,
    ptrdiff_t k_cache_strides_0,
    ptrdiff_t k_cache_strides_1,
    ptrdiff_t k_cache_strides_2,
    ptrdiff_t k_cache_strides_3,
    ptrdiff_t v_cache_strides_0,
    ptrdiff_t v_cache_strides_1,
    ptrdiff_t v_cache_strides_2,
    ptrdiff_t v_cache_strides_3,
    ptrdiff_t k_strides_0,
    ptrdiff_t k_strides_1,
    ptrdiff_t k_strides_2,
    ptrdiff_t k_strides_3,
    ptrdiff_t v_strides_0,
    ptrdiff_t v_strides_1,
    ptrdiff_t v_strides_2,
    ptrdiff_t v_strides_3) {
    kvCachingKernel<Tdata, Tidx>(k_cache, v_cache, k, v, past_kv_lengths,
                                 batch_size, num_kv_heads, max_seq_len, seq_len, hidden_dim,
                                 k_cache_strides_0, k_cache_strides_1, k_cache_strides_2, k_cache_strides_3,
                                 v_cache_strides_0, v_cache_strides_1, v_cache_strides_2, v_cache_strides_3,
                                 k_strides_0, k_strides_1, k_strides_2, k_strides_3,
                                 v_strides_0, v_strides_1, v_strides_2, v_strides_3);
}

namespace op::kv_caching::metax {

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

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

infiniStatus_t Descriptor::create(
    infiniopHandle_t handle,
    Descriptor **desc_ptr,
    infiniopTensorDescriptor_t k_cache,
    infiniopTensorDescriptor_t v_cache,
    infiniopTensorDescriptor_t k,
    infiniopTensorDescriptor_t v,
    infiniopTensorDescriptor_t past_kv_lengths) {
    auto info = KVCachingInfo::createKVCachingInfo(k_cache, v_cache, k, v, past_kv_lengths);
    CHECK_RESULT(info);

    *desc_ptr = new Descriptor(
        new Opaque{reinterpret_cast<device::metax::Handle *>(handle)->internal()},
        info.take(), 0, handle->device, handle->device_id);
    return INFINI_STATUS_SUCCESS;
}

template <unsigned int BLOCK_SIZE, typename Tdata, typename Tidx>
infiniStatus_t launchKernel(const KVCachingInfo &info,
                            Tdata *k_cache,
                            Tdata *v_cache,
                            const Tdata *k,
                            const Tdata *v,
                            const Tidx *past_kv_lengths,
                            hcStream_t stream, void *workspace) {

    int batch_size = static_cast<int>(info.batch_size);
    int num_kv_heads = static_cast<int>(info.num_kv_heads);
    int max_seq_len = static_cast<int>(info.max_seq_len);
    int hidden_dim = static_cast<int>(info.hidden_dim);

    int seq_len = static_cast<int>(info.seq_len);

    int total = batch_size * num_kv_heads * seq_len * hidden_dim;

    ptrdiff_t k_cache_strides_0 = info.k_cache_strides_0;
    ptrdiff_t k_cache_strides_1 = info.k_cache_strides_1;
    ptrdiff_t k_cache_strides_2 = info.k_cache_strides_2;
    ptrdiff_t k_cache_strides_3 = info.k_cache_strides_3;

    ptrdiff_t v_cache_strides_0 = info.v_cache_strides_0;
    ptrdiff_t v_cache_strides_1 = info.v_cache_strides_1;
    ptrdiff_t v_cache_strides_2 = info.v_cache_strides_2;
    ptrdiff_t v_cache_strides_3 = info.v_cache_strides_3;

    ptrdiff_t k_strides_0 = info.k_strides_0;
    ptrdiff_t k_strides_1 = info.k_strides_1;
    ptrdiff_t k_strides_2 = info.k_strides_2;
    ptrdiff_t k_strides_3 = info.k_strides_3;

    ptrdiff_t v_strides_0 = info.v_strides_0;
    ptrdiff_t v_strides_1 = info.v_strides_1;
    ptrdiff_t v_strides_2 = info.v_strides_2;
    ptrdiff_t v_strides_3 = info.v_strides_3;

    int num_blocks = (total + BLOCK_SIZE - 1) / BLOCK_SIZE;

    kvCaching<Tdata, Tidx>
        <<<num_blocks, BLOCK_SIZE, 0, stream>>>(k_cache, v_cache, k, v, past_kv_lengths,
                                                batch_size, num_kv_heads, max_seq_len, seq_len, hidden_dim,
                                                k_cache_strides_0, k_cache_strides_1, k_cache_strides_2, k_cache_strides_3,
                                                v_cache_strides_0, v_cache_strides_1, v_cache_strides_2, v_cache_strides_3,
                                                k_strides_0, k_strides_1, k_strides_2, k_strides_3,
                                                v_strides_0, v_strides_1, v_strides_2, v_strides_3);
    return INFINI_STATUS_SUCCESS;
}

infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size,
                                     void *k_cache,
                                     void *v_cache,
                                     const void *k,
                                     const void *v,
                                     const void *past_kv_lengths,
                                     void *stream_) const {
    hcStream_t stream = (hcStream_t)stream_;

#define LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DTYPE(BLOCK_SIZE, TDATA, TIDX) \
    launchKernel<BLOCK_SIZE, TDATA, TIDX>(_info, (TDATA *)k_cache, (TDATA *)v_cache, \
                                          (const TDATA *)k, (const TDATA *)v, \
                                          (const TIDX *)past_kv_lengths, stream, workspace)

#define LAUNCH_KERNEL_WITH_BLOCK_SIZE(BLOCK_SIZE, TDATA) \
    if (_info.past_len_dtype == INFINI_DTYPE_I32) { \
        return LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DTYPE(BLOCK_SIZE, TDATA, int32_t); \
    } else { /* INFINI_DTYPE_I64 */ \
        return LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DTYPE(BLOCK_SIZE, TDATA, int64_t); \
    }

#define LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DATA_TYPE(BLOCK_SIZE) \
    { \
        if (_info.dtype == INFINI_DTYPE_F16) { \
            LAUNCH_KERNEL_WITH_BLOCK_SIZE(BLOCK_SIZE, half) \
        } else if (_info.dtype == INFINI_DTYPE_F32) { \
            LAUNCH_KERNEL_WITH_BLOCK_SIZE(BLOCK_SIZE, float) \
        } else if (_info.dtype == INFINI_DTYPE_BF16) { \
            LAUNCH_KERNEL_WITH_BLOCK_SIZE(BLOCK_SIZE, __hpcc_bfloat16) \
        } else { \
            return INFINI_STATUS_BAD_TENSOR_DTYPE; \
        } \
    }

    // Choose block size based on device capabilities
    if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_1024) {
        LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DATA_TYPE(METAX_BLOCK_SIZE_1024)
    } else if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_512) {
        LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DATA_TYPE(METAX_BLOCK_SIZE_512)
    } else if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_2048) {
        LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DATA_TYPE(METAX_BLOCK_SIZE_2048)
    } else if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_4096) {
        LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DATA_TYPE(METAX_BLOCK_SIZE_4096)
    } else {
        return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED;
    }
    return INFINI_STATUS_SUCCESS;
}

} // namespace op::kv_caching::metax
