// moe_ops_impl.cu — Implement the 5 missing MoE functions
//
// These functions are declared in ixformer.h (from xllm upstream)
// but NOT present in the base image's libixformer.so.
//
// We implement them using available primitives:
//   - cuinferCustomGemm (from libcuinfer.so) for group_gemm
//   - Pure CUDA kernels for topk_softmax, moe_compute_index, expand, combine
//   - ixformer::functions::cuinfer_gemm (from libixformer.so) as fallback
//
// Reference AST chain:
//   xllm/core/kernels/ilu/fused_moe.cpp  → calls these 5 functions
//   xllm/core/kernels/ilu/group_gemm.cpp → calls moe_w16a16_group_gemm
//   xllm/core/kernels/ilu/ixformer.h     → declares them in ixformer::infer
//
// We provide them in the SAME namespace so ix_full_bridge_v2.cpp links cleanly.

#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <optional>
#include <vector>
#include <numeric>

// ============================================================================
// Forward-declare cuinfer C API (from libcuinfer.so, confirmed in symbol dump)
// ============================================================================
extern "C" {

typedef struct cuinferContext* cuinferHandle_t;
typedef enum { CUINFER_STATUS_SUCCESS = 0 } cuinferStatus_t;
typedef enum {
    CUINFER_OP_TENSOR_OP_N = 0,
    CUINFER_OP_TENSOR_OP_T = 1,
} cuinferOperation_t;
typedef enum {
    CUINFER_GEMM_DEFAULT = 0,
} cuinferGEMMCustomOption_t;
typedef enum {
    CUINFER_POINTER_MODE_HOST = 0,
} cuinferPointerMode_t;

cuinferStatus_t cuinferCreate(cuinferHandle_t* handle);
cuinferStatus_t cuinferDestroy(cuinferHandle_t handle);
cuinferStatus_t cuinferSetStream(cuinferHandle_t handle, cudaStream_t stream);

cuinferStatus_t cuinferCustomGemm(
    cuinferHandle_t handle, cudaStream_t stream,
    cuinferPointerMode_t ptrMode,
    cuinferOperation_t transa, cuinferOperation_t transb,
    int m, int n, int k,
    const void* alpha,
    const void* A, cudaDataType_t Atype, int lda, long long int strideA,
    const void* B, cudaDataType_t Btype, int ldb, long long int strideB,
    const void* beta,
    void* C, cudaDataType_t Ctype, int ldc, long long int strideC,
    int batchCount,
    cudaDataType_t computeType, cudaDataType_t scaleType,
    const void* customHostPtr, const void* customDevicePtr,
    cuinferGEMMCustomOption_t customOption);

}  // extern "C"


// ============================================================================
// Kernel 1: topk_softmax
// Adapted from moe_topk_softmax_v3.cu (already working, 64-expert specialized)
// ============================================================================

// Qwen3.5-27B: 128 routed experts
// Block size = 128 threads (1 thread per expert for ≤128 experts)
static constexpr int MOE_MAX_EXPERTS = 128;
static constexpr int MOE_BLOCK = 128;

// All reductions use blockDim.x (dynamic block size, power-of-2)
__device__ float smem_reduce_max(float val, float* smem) {
    int tid = threadIdx.x;
    smem[tid] = val;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) smem[tid] = fmaxf(smem[tid], smem[tid + s]);
        __syncthreads();
    }
    return smem[0];
}

__device__ float smem_reduce_sum(float val, float* smem) {
    int tid = threadIdx.x;
    smem[tid] = val;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }
    return smem[0];
}

__device__ void smem_argmax(float val, int idx, float* s_val, int* s_idx) {
    int tid = threadIdx.x;
    s_val[tid] = val;
    s_idx[tid] = idx;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s && s_val[tid + s] > s_val[tid]) {
            s_val[tid] = s_val[tid + s];
            s_idx[tid] = s_idx[tid + s];
        }
        __syncthreads();
    }
}

__global__ void topk_softmax_kernel(
    const float* __restrict__ input,
    float* __restrict__ topk_weights,
    int32_t* __restrict__ topk_indices,
    int32_t* __restrict__ token_expert_indices,
    int num_tokens, int num_experts, int topk, bool renormalize
) {
    int row = blockIdx.x;
    if (row >= num_tokens) return;
    int tid = threadIdx.x;

    extern __shared__ char shared_buf[];
    float* smem = (float*)shared_buf;
    int* smem_idx = (int*)(smem + blockDim.x);

    // num_experts passed via gridDim.y (encoded), or read from shared
    // We use a separate parameter for clarity
    float val = (tid < num_experts) ? input[row * num_experts + tid] : -1e30f;

    // Softmax
    float row_max = smem_reduce_max(val, smem);
    val = (tid < num_experts) ? expf(val - row_max) : 0.0f;
    float row_sum = smem_reduce_sum(val, smem);
    val *= (1.0f / row_sum);

    float* out_w = topk_weights + row * topk;
    int32_t* out_idx = topk_indices + row * topk;
    int32_t* out_src = token_expert_indices + row * topk;

    float my_val = val;
    float topk_sum = 0.0f;

    for (int ki = 0; ki < topk; ki++) {
        smem_argmax(my_val, tid, smem, smem_idx);
        float winner_val = smem[0];
        int winner_idx = smem_idx[0];
        __syncthreads();

        if (tid == 0) {
            out_w[ki] = winner_val;
            out_idx[ki] = winner_idx;
            out_src[ki] = row;
        }
        topk_sum += winner_val;
        if (tid == winner_idx) my_val = -1.0f;
        __syncthreads();
    }

    if (renormalize && tid == 0) {
        float inv = 1.0f / (topk_sum + 1e-8f);
        for (int ki = 0; ki < topk; ki++)
            out_w[ki] *= inv;
    }
}


// ============================================================================
// Kernel 2: moe_compute_token_index
// Histogram + prefix sum + scatter — from xllm_kernels/cuda/moe_compute_index.cu
// ============================================================================

__global__ void histogram_kernel(
    const int32_t* __restrict__ expert_ids,
    int32_t* __restrict__ expert_sizes,
    int num_elements, int num_experts
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < num_elements) {
        int eid = expert_ids[idx];
        if (eid >= 0 && eid < num_experts) {
            atomicAdd(&expert_sizes[eid], 1);
        }
    }
}

__global__ void place_indices_kernel(
    const int32_t* __restrict__ expert_ids,
    int32_t* __restrict__ expert_offsets,  // will be atomicAdd'd
    int32_t* __restrict__ src_dst,
    int32_t* __restrict__ dst_src,
    int num_elements
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < num_elements) {
        int eid = expert_ids[idx];
        int pos = atomicAdd(&expert_offsets[eid], 1);
        src_dst[idx] = pos;  // where token idx goes in sorted order
        dst_src[pos] = idx;  // reverse mapping
    }
}


// ============================================================================
// Kernel 3: moe_expand_input
// Gather-based expand: output[i] = input[gather_index[i]]
// ============================================================================

template <typename scalar_t>
__global__ void expand_input_kernel(
    scalar_t* __restrict__ output,
    const scalar_t* __restrict__ input,
    const int32_t* __restrict__ dst_to_src,
    int num_output_tokens, int hidden_size
) {
    int token = blockIdx.x;
    if (token >= num_output_tokens) return;

    int src_token = dst_to_src[token];
    const scalar_t* src = input + (int64_t)src_token * hidden_size;
    scalar_t* dst = output + (int64_t)token * hidden_size;

    for (int h = threadIdx.x; h < hidden_size; h += blockDim.x) {
        dst[h] = src[h];
    }
}


// ============================================================================
// Kernel 4: moe_combine_result (weighted sum of expert outputs)
// output[t] = sum_k( weight[t][k] * gemm2_output[flat_index(t,k)] )
// ============================================================================

template <typename scalar_t>
__global__ void combine_result_kernel(
    scalar_t* __restrict__ output,      // [N, H]
    const scalar_t* __restrict__ input,  // [N*topk, H]
    const float* __restrict__ weights,   // [N, topk]
    int num_tokens, int topk, int hidden_size
) {
    int token = blockIdx.x;
    if (token >= num_tokens) return;

    for (int h = threadIdx.x; h < hidden_size; h += blockDim.x) {
        float acc = 0.0f;
        for (int k = 0; k < topk; k++) {
            int flat = token * topk + k;
            float w = weights[token * topk + k];
            acc += w * __half2float(input[flat * hidden_size + h]);
        }
        output[token * hidden_size + h] = __float2half(acc);
    }
}

// Float specialization
template <>
__global__ void combine_result_kernel<float>(
    float* __restrict__ output,
    const float* __restrict__ input,
    const float* __restrict__ weights,
    int num_tokens, int topk, int hidden_size
) {
    int token = blockIdx.x;
    if (token >= num_tokens) return;

    for (int h = threadIdx.x; h < hidden_size; h += blockDim.x) {
        float acc = 0.0f;
        for (int k = 0; k < topk; k++) {
            int flat = token * topk + k;
            float w = weights[token * topk + k];
            acc += w * input[flat * hidden_size + h];
        }
        output[token * hidden_size + h] = acc;
    }
}


// ============================================================================
// C++ wrapper functions — ixformer::infer namespace
// These provide the MISSING symbols that ix_full_bridge_v2.cpp needs.
// ============================================================================

namespace ixformer { namespace infer {

void topk_softmax(
    torch::Tensor& topk_weights,
    torch::Tensor& topk_indices,
    torch::Tensor& token_expert_indices,
    torch::Tensor& gating_output,
    bool renormalize
) {
    int num_tokens = gating_output.size(0);
    int num_experts = gating_output.size(1);
    int topk = topk_weights.size(1);
    auto stream = c10::cuda::getCurrentCUDAStream();

    auto input_f32 = gating_output.to(torch::kFloat32).contiguous();

    // Block size must be >= num_experts, round up to next power of 2
    int block_size = 1;
    while (block_size < num_experts) block_size <<= 1;
    TORCH_CHECK(block_size <= 1024, "Too many experts for topk kernel: ", num_experts);

    size_t smem_bytes = block_size * (sizeof(float) + sizeof(int));
    topk_softmax_kernel<<<num_tokens, block_size, smem_bytes, stream>>>(
        input_f32.data_ptr<float>(),
        topk_weights.data_ptr<float>(),
        topk_indices.data_ptr<int32_t>(),
        token_expert_indices.data_ptr<int32_t>(),
        num_tokens, num_experts, topk, renormalize);
}

void moe_compute_token_index_api(
    torch::Tensor& topk_ids,
    torch::Tensor& src_dst,
    torch::Tensor& dst_src,
    torch::Tensor& expert_sizes_gpu,
    const std::optional<torch::Tensor>& expert_mask,
    const std::optional<torch::Tensor>& expert_sizes_cpu,
    const std::optional<torch::Tensor>& expand_tokens_gpu,
    int64_t start_expert_id,
    int64_t end_expert_id,
    int64_t num_experts
) {
    auto stream = c10::cuda::getCurrentCUDAStream();
    int num_elements = topk_ids.numel();

    // Zero expert_sizes
    cudaMemsetAsync(expert_sizes_gpu.data_ptr<int32_t>(), 0,
                    num_experts * sizeof(int32_t), stream);

    // Phase 1: histogram
    int blocks1 = (num_elements + 255) / 256;
    histogram_kernel<<<blocks1, 256, 0, stream>>>(
        topk_ids.data_ptr<int32_t>(),
        expert_sizes_gpu.data_ptr<int32_t>(),
        num_elements, num_experts);

    // Phase 2: prefix sum for offsets (exclusive scan on GPU)
    // Use a separate buffer for offsets, then reset for place_indices
    auto expert_offsets = torch::zeros({num_experts}, topk_ids.options().dtype(torch::kInt32));
    // Copy sizes → do exclusive scan on CPU (small: 64 experts)
    auto sizes_cpu = expert_sizes_gpu.to(torch::kCPU);
    auto offsets_cpu = torch::zeros({num_experts}, torch::dtype(torch::kInt32));
    int32_t* s = sizes_cpu.data_ptr<int32_t>();
    int32_t* o = offsets_cpu.data_ptr<int32_t>();
    int32_t running = 0;
    for (int i = 0; i < num_experts; i++) {
        o[i] = running;
        running += s[i];
    }
    expert_offsets = offsets_cpu.to(topk_ids.device());

    // Phase 3: place indices
    int blocks3 = (num_elements + 255) / 256;
    place_indices_kernel<<<blocks3, 256, 0, stream>>>(
        topk_ids.data_ptr<int32_t>(),
        expert_offsets.data_ptr<int32_t>(),
        src_dst.data_ptr<int32_t>(),
        dst_src.data_ptr<int32_t>(),
        num_elements);
}

void moe_expand_input(
    torch::Tensor outputs,
    torch::Tensor inputs,
    torch::Tensor dst_to_src,
    const std::optional<torch::Tensor>& src_to_dst,
    int64_t dst_tokens,
    int64_t expand_factor
) {
    auto stream = c10::cuda::getCurrentCUDAStream();
    int hidden_size = inputs.size(1);
    int block = std::min(hidden_size, 256);

    AT_DISPATCH_FLOATING_TYPES_AND_HALF(inputs.scalar_type(), "expand_input", [&] {
        expand_input_kernel<scalar_t><<<dst_tokens, block, 0, stream>>>(
            outputs.data_ptr<scalar_t>(),
            inputs.data_ptr<scalar_t>(),
            dst_to_src.data_ptr<int32_t>(),
            dst_tokens, hidden_size);
    });
}

void moe_w16a16_group_gemm(
    torch::Tensor output,
    torch::Tensor inputs,
    torch::Tensor weights,
    torch::Tensor tokens_per_experts,
    const std::optional<torch::Tensor>& dst_to_src,
    const std::optional<torch::Tensor>& bias,
    std::string format,
    int64_t persistent,
    int64_t output_n
) {
    // Implementation: loop over experts, call cuinferCustomGemm for each
    // weights: [num_experts, N, K] with format "TN" means transB
    // For each expert e with count tokens:
    //   A = inputs[offset:offset+count, :]        (count × K, row-major)
    //   B = weights[e, :, :]                      (N × K, needs transB)
    //   C = output[offset:offset+count, :]        (count × N, row-major)
    //   GEMM: C = A × B^T  → (count, K) × (K, N) = (count, N)

    auto stream = c10::cuda::getCurrentCUDAStream();
    int num_experts = weights.size(0);
    int N = weights.size(1);  // output dim
    int K = weights.size(2);  // input dim

    // Get token counts on CPU
    auto counts_cpu = tokens_per_experts.to(torch::kCPU).to(torch::kInt32);
    int32_t* counts = counts_cpu.data_ptr<int32_t>();

    // Create cuinfer handle
    cuinferHandle_t handle;
    cuinferCreate(&handle);
    cuinferSetStream(handle, stream);

    float alpha = 1.0f, beta = 0.0f;

    int offset = 0;
    for (int e = 0; e < num_experts; e++) {
        int M = counts[e];
        if (M <= 0) continue;

        // A: inputs[offset : offset+M, :]  → M × K
        // B: weights[e, :, :]               → N × K (transposed: compute A × B^T)
        // C: output[offset : offset+M, :]   → M × N
        const void* A_ptr = (const char*)inputs.data_ptr() +
            (int64_t)offset * K * inputs.element_size();
        const void* B_ptr = (const char*)weights.data_ptr() +
            (int64_t)e * N * K * weights.element_size();
        void* C_ptr = (char*)output.data_ptr() +
            (int64_t)offset * N * output.element_size();

        cudaDataType_t dtype = (inputs.scalar_type() == torch::kFloat16)
            ? CUDA_R_16F : CUDA_R_32F;

        // cuinferCustomGemm: row-major convention
        // We want C = A × B^T
        // In cuinfer (column-major internally): transa=N, transb=T
        // M_gemm = M (rows of C), N_gemm = N (cols of C), K_gemm = K
        cuinferCustomGemm(
            handle, stream,
            CUINFER_POINTER_MODE_HOST,
            CUINFER_OP_TENSOR_OP_N,  // transa = no transpose
            CUINFER_OP_TENSOR_OP_T,  // transb = transpose (TN format)
            M, N, K,
            &alpha,
            A_ptr, dtype, K, 0,       // lda=K for row-major A
            B_ptr, dtype, K, 0,       // ldb=K for row-major B (will be transposed)
            &beta,
            C_ptr, dtype, N, 0,       // ldc=N for row-major C
            1,                         // batchCount=1
            CUDA_R_32F,               // computeType
            CUDA_R_32F,               // scaleType
            nullptr, nullptr,          // custom pointers
            CUINFER_GEMM_DEFAULT);

        offset += M;
    }

    cuinferDestroy(handle);
}

void moe_output_reduce_sum(
    torch::Tensor outputs,
    torch::Tensor inputs,
    const std::optional<torch::Tensor>& mul_weight,
    const std::optional<torch::Tensor>& mask,
    const std::optional<torch::Tensor>& extra_residual,
    double scaling_factor
) {
    // inputs: [N, topk, H] — expert outputs per token
    // mul_weight: [N, topk]  — router weights
    // outputs: [N, H]       — weighted sum
    auto stream = c10::cuda::getCurrentCUDAStream();
    int num_tokens = inputs.size(0);
    int topk = inputs.size(1);
    int hidden_size = inputs.size(2);
    int block = std::min(hidden_size, 256);

    // Reshape inputs to [N*topk, H] for the kernel
    auto input_flat = inputs.reshape({num_tokens * topk, hidden_size});

    if (inputs.scalar_type() == torch::kFloat16) {
        combine_result_kernel<__half><<<num_tokens, block, 0, stream>>>(
            reinterpret_cast<__half*>(outputs.data_ptr()),
            reinterpret_cast<const __half*>(input_flat.data_ptr()),
            mul_weight.value().data_ptr<float>(),
            num_tokens, topk, hidden_size);
    } else {
        combine_result_kernel<float><<<num_tokens, block, 0, stream>>>(
            outputs.data_ptr<float>(),
            input_flat.data_ptr<float>(),
            mul_weight.value().data_ptr<float>(),
            num_tokens, topk, hidden_size);
    }
}

}}  // namespace ixformer::infer