/**
 * Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * See file LICENSE for terms.
 */

#include "reduce_scatterv/reduce_scatterv.h"
#include <ucc/api/ucc.h>

#include "core/ucc_ee.h"
#include "utils/arch/cuda_def.h"
#include "tl_cuda_nvls.h"
#include "kernels/reduce_scatter_kernel.h"

#include "components/ec/ucc_ec.h"
#include "components/ec/cuda/ec_cuda_resources.h"

ucc_status_t ucc_tl_cuda_reduce_scatterv_nvls_start(ucc_coll_task_t *coll_task)
{
    ucc_tl_cuda_task_t *task = ucc_derived_of(coll_task, ucc_tl_cuda_task_t);
    ucc_tl_cuda_team_t *team = TASK_TEAM(task);
    ucc_coll_args_t    *args = &TASK_ARGS(task);
    ucc_datatype_t      dt   = task->reduce_scatterv_nvls.dt;
    ucc_ee_h            ee   = coll_task->ee;
    cudaStream_t stream = (ee) ? (cudaStream_t)ee->ee_context : team->stream;
    ucc_ec_cuda_event_t *ec_event = (ucc_ec_cuda_event_t *)task
                                        ->reduce_scatterv_nvls.evt_completion;
    cudaEvent_t  evt      = ec_event->event;
    uint32_t     sm_count = UCC_TL_CUDA_TEAM_LIB(team)->cfg.nvls_sm_count;
    uint32_t     threads  = UCC_TL_CUDA_TEAM_LIB(team)->cfg.nvls_threads;
    CUdeviceptr  mc_va    = task->reduce_scatterv_nvls.mc_va;
    CUdeviceptr  uc_va    = task->reduce_scatterv_nvls.uc_va;
    ucc_status_t status;
    size_t       src_size_bytes;
    void        *sbuf;
    void        *rbuf;

    if (args->coll_type == UCC_COLL_TYPE_REDUCE_SCATTER) {
        if (UCC_IS_INPLACE(*args)) {
            sbuf = args->dst.info.buffer;
            rbuf = PTR_OFFSET(
                sbuf, task->reduce_scatterv_nvls.offset * sizeof(uint32_t));
            src_size_bytes = args->dst.info.count * UCC_TL_TEAM_SIZE(team) *
                             ucc_dt_size(dt);
        } else {
            sbuf           = args->src.info.buffer;
            rbuf           = args->dst.info.buffer;
            src_size_bytes = args->src.info.count * ucc_dt_size(dt);
        }
    } else {
        /* RSV: reduce_scatterv */
        if (UCC_IS_INPLACE(*args)) {
            sbuf = args->dst.info_v.buffer;
            rbuf = PTR_OFFSET(
                sbuf, task->reduce_scatterv_nvls.offset * sizeof(uint32_t));
            src_size_bytes = ucc_coll_args_get_total_count(
                                 args,
                                 args->dst.info_v.counts,
                                 UCC_TL_TEAM_SIZE(team)) *
                             ucc_dt_size(dt);
        } else {
            sbuf           = args->src.info.buffer;
            rbuf           = args->dst.info_v.buffer;
            src_size_bytes = args->src.info.count * ucc_dt_size(dt);
        }
    }

    /* copy src buffer to symmetric memory first */
    status = CUDA_FUNC(cudaMemcpyAsync(
        (void *)uc_va, sbuf, src_size_bytes, cudaMemcpyDeviceToDevice, stream));
    if (ucc_unlikely(status != UCC_OK)) {
        return status;
    }

    status = post_reduce_scatter_kernel(
        stream,
        sm_count,
        threads,
        (CUdeviceptr)rbuf,
        mc_va,
        TASK_NVLS_CONTROL_MC(task),
        TASK_NVLS_CONTROL_UC(task),
        task->reduce_scatterv_nvls.coll_id,
        task->reduce_scatterv_nvls.offset,
        task->reduce_scatterv_nvls.count,
        dt,
        UCC_TL_TEAM_SIZE(team));
    if (ucc_unlikely(status != UCC_OK)) {
        tl_error(UCC_TASK_LIB(task), "failed to post reduce scatter kernel");
        return status;
    }
    status = CUDA_FUNC(cudaEventRecord(evt, stream));
    if (ucc_unlikely(status != UCC_OK)) {
        return status;
    }

    return ucc_progress_queue_enqueue(UCC_TL_CORE_CTX(team)->pq, &task->super);
}

void ucc_tl_cuda_reduce_scatterv_nvls_progress(ucc_coll_task_t *coll_task)
{
    ucc_tl_cuda_task_t  *task = ucc_derived_of(coll_task, ucc_tl_cuda_task_t);
    ucc_ec_cuda_event_t *ec_event = (ucc_ec_cuda_event_t *)task
                                        ->reduce_scatterv_nvls.evt_completion;
    cudaEvent_t evt    = ec_event->event;

    task->super.status = cuda_error_to_ucc_status(cudaEventQuery(evt));
}

ucc_status_t ucc_tl_cuda_reduce_scatterv_nvls_triggered_post(
    ucc_ee_h ee, ucc_ev_t *ev, ucc_coll_task_t *coll_task) // NOLINT
{
    ucc_tl_cuda_task_t *task = ucc_derived_of(coll_task, ucc_tl_cuda_task_t);
    ucc_status_t        status;
    ucc_ev_t            post_event;

    ucc_assert(ee->ee_type == UCC_EE_CUDA_STREAM);
    coll_task->ee = ee;
    tl_trace(UCC_TASK_LIB(task), "triggered post. task:%p", coll_task);

    status = coll_task->post(coll_task);
    if (ucc_likely(status == UCC_OK)) {
        post_event.ev_type         = UCC_EVENT_COLLECTIVE_POST;
        post_event.ev_context_size = 0;
        post_event.ev_context      = NULL;
        post_event.req             = &coll_task->super;
        status                     = ucc_ee_set_event_internal(
            coll_task->ee, &post_event, &coll_task->ee->event_out_queue);
        if (ucc_unlikely(status != UCC_OK)) {
            tl_error(
                UCC_TASK_LIB(task),
                "failed to set EE event: %s",
                ucc_status_string(status));
            return status;
        }
    }
    return status;
}

ucc_status_t ucc_tl_cuda_reduce_scatterv_nvls_finalize(ucc_coll_task_t *task)
{
    ucc_tl_cuda_task_t *tl_task = ucc_derived_of(task, ucc_tl_cuda_task_t);
    tl_trace(
        UCC_TASK_LIB(tl_task), "task: %p reduce_scatterv_nvls_finalize", task);

    ucc_ec_destroy_event(
        tl_task->reduce_scatterv_nvls.evt_completion, UCC_EE_CUDA_STREAM);

    ucc_tl_cuda_task_put(tl_task);
    return UCC_OK;
}

ucc_status_t ucc_tl_cuda_reduce_scatterv_nvls_init_common(
    ucc_base_coll_args_t *coll_args, ucc_base_team_t *tl_team,
    ucc_coll_task_t **task_p, nvls_get_count_fn_t get_offset,
    nvls_get_count_fn_t get_count)
{
    ucc_tl_cuda_team_t *team  = ucc_derived_of(tl_team, ucc_tl_cuda_team_t);
    ucc_rank_t          trank = UCC_TL_TEAM_RANK(team);
    ucc_datatype_t      dt    = (coll_args->args.coll_type ==
                         UCC_COLL_TYPE_REDUCE_SCATTER)
                                    ? coll_args->args.dst.info.datatype
                                    : coll_args->args.dst.info_v.datatype;
    ucc_tl_cuda_task_t *task;
    ucc_status_t        status;
    size_t              offset_elements;
    size_t              count_elements;
    size_t              offset;
    size_t              count;

    /* Validate op and datatype before allocating task */
    if (coll_args->args.op != UCC_OP_SUM) {
        tl_debug(
            UCC_TL_TEAM_LIB(team),
            "NVLS reduce scatter(v) supported only with SUM operation");
        return UCC_ERR_NOT_SUPPORTED;
    }
    if (dt != UCC_DT_FLOAT32 && dt != UCC_DT_BFLOAT16 && dt != UCC_DT_INT32 &&
        dt != UCC_DT_UINT32 && dt != UCC_DT_INT64 && dt != UCC_DT_UINT64) {
        tl_debug(
            UCC_TL_TEAM_LIB(team),
            "NVLS reduce scatter(v) supported only with "
            "float32/bfloat16/int32/uint32/int64/uint64");
        return UCC_ERR_NOT_SUPPORTED;
    }

    status = ucc_tl_cuda_task_init(coll_args, team, &task);
    if (ucc_unlikely(status != UCC_OK)) {
        return status;
    }

    offset_elements = get_offset(task, trank);
    count_elements  = get_count(task, trank);

    /* Convert from datatype elements to uint32_t indices for the kernel.
     * For float32/int32/uint32: 1 element = 1 uint32_t
     * For int64/uint64: 1 element = 2 uint32_t
     * For bfloat16: 2 elements = 1 uint32_t */
    if (dt == UCC_DT_FLOAT32 || dt == UCC_DT_INT32 || dt == UCC_DT_UINT32) {
        offset = offset_elements;
        count  = count_elements;
    } else if (dt == UCC_DT_INT64 || dt == UCC_DT_UINT64) {
        offset = offset_elements * 2;
        count  = count_elements * 2;
    } else { /* UCC_DT_BFLOAT16 */
        if (offset_elements % 2 != 0 || count_elements % 2 != 0) {
            tl_debug(
                UCC_TL_TEAM_LIB(team),
                "BF16 offset and count must be even, got offset=%zu "
                "count=%zu",
                offset_elements,
                count_elements);
            goto err_task_put;
        }
        offset = offset_elements / 2;
        count  = count_elements / 2;
    }

    /* NVLS requires 16-byte alignment (4 uint32_t elements) */
    if (ucc_unlikely(offset % 4 != 0)) {
        tl_debug(
            UCC_TL_TEAM_LIB(team),
            "NVLS requires 16-byte alignment for offset, got offset=%zu",
            offset);
        goto err_task_put;
    }
    if (ucc_unlikely(count % 4 != 0)) {
        tl_debug(
            UCC_TL_TEAM_LIB(team),
            "NVLS requires 16-byte alignment for count, got count=%zu",
            count);
        goto err_task_put;
    }

    status = ucc_ec_create_event(
        &task->reduce_scatterv_nvls.evt_completion, UCC_EE_CUDA_STREAM);
    if (ucc_unlikely(status != UCC_OK)) {
        tl_error(UCC_TL_TEAM_LIB(team), "failed to create CUDA event");
        goto err_task_put;
    }

    task->reduce_scatterv_nvls.dt      = dt;
    task->reduce_scatterv_nvls.offset  = offset;
    task->reduce_scatterv_nvls.count   = count;
    task->reduce_scatterv_nvls.mc_va   = (CUdeviceptr)TASK_SYMMETRIC_MC(task);
    task->reduce_scatterv_nvls.uc_va   = (CUdeviceptr)TASK_SYMMETRIC_UC(task);
    task->reduce_scatterv_nvls.coll_id = team->nvls.coll_ids[task->coll_id]++;

    task->super.post                   = ucc_tl_cuda_reduce_scatterv_nvls_start;
    task->super
        .triggered_post  = ucc_tl_cuda_reduce_scatterv_nvls_triggered_post;
    task->super.progress = ucc_tl_cuda_reduce_scatterv_nvls_progress;
    task->super.finalize = ucc_tl_cuda_reduce_scatterv_nvls_finalize;

    *task_p              = &task->super;
    return UCC_OK;

err_task_put:
    ucc_tl_cuda_task_put(task);
    return UCC_ERR_NOT_SUPPORTED;
}

ucc_status_t ucc_tl_cuda_reduce_scatterv_nvls_init(
    ucc_base_coll_args_t *coll_args, ucc_base_team_t *tl_team,
    ucc_coll_task_t **task_p)
{
    return ucc_tl_cuda_reduce_scatterv_nvls_init_common(
        coll_args,
        tl_team,
        task_p,
        ucc_tl_cuda_reduce_scatterv_get_offset,
        ucc_tl_cuda_reduce_scatterv_get_count);
}
