#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <mpi.h>
#include "kupl_shm.h"  // KUPL库头文件（来自鲲鹏HPCKit）

// 全局定义缓冲区大小（可修改，用于测试不同消息大小的性能）
#define BUF_SIZE_SMALL 4    // 小消息（4个int，16字节）
#define BUF_SIZE_MEDIUM 1024 // 中消息（1024个int，4KB）
#define BUF_SIZE_LARGE 1048576 // 大消息（1MB）

// 1. 实现oob_allgather回调函数（进程地址收集）
static int oob_allgather_callback(const void *sendbuf, void *recvbuf, int size, void *group, kupl_shm_datatype_t datatype) {
    int ret;
    // 判断数据类型，确保与MPI数据类型匹配
    switch (datatype) {
        case KUPL_SHM_DATATYPE_CHAR:
            ret = MPI_Allgather(sendbuf, size, MPI_CHAR, recvbuf, size, MPI_CHAR, (MPI_Comm)group);
            break;
        case KUPL_SHM_DATATYPE_INT:
            ret = MPI_Allgather(sendbuf, size, MPI_INT, recvbuf, size, MPI_INT, (MPI_Comm)group);
            break;
        case KUPL_SHM_DATATYPE_DOUBLE:
            ret = MPI_Allgather(sendbuf, size, MPI_DOUBLE, recvbuf, size, MPI_DOUBLE, (MPI_Comm)group);
            break;
        default:
            fprintf(stderr, "Error: Unsupported datatype in oob_allgather_callback\n");
            return -1; // 不支持的数据类型，返回失败
    }
    // 检查MPI_Allgather调用结果
    if (ret != MPI_SUCCESS) {
        fprintf(stderr, "Error: MPI_Allgather failed in oob_allgather_callback\n");
        return -1;
    }
    return 0; // 执行成功
}

// 2. 实现oob_barrier回调函数（进程同步）
static int oob_barrier_callback(void *group) {
    int ret = MPI_Barrier((MPI_Comm)group);
    // 检查MPI_Barrier调用结果
    if (ret != MPI_SUCCESS) {
        fprintf(stderr, "Error: MPI_Barrier failed in oob_barrier_callback\n");
        return -1;
    }
    return 0; // 执行成功
}

// 3. 数据校验函数（验证接收的数据与发送的数据是否一致）
static int data_verify(int *send_buf, int *recv_buf, int buf_size) {
    for (int i = 0; i < buf_size; i++) {
        if (send_buf[i] != recv_buf[i]) {
            fprintf(stderr, "Data verify failed: send_buf[%d] = %d, recv_buf[%d] = %d\n", i, send_buf[i], i, recv_buf[i]);
            return -1;
        }
    }
    return 0; // 校验成功
}

int main(int argc, char **argv) {
    int ret;
    int world_rank, world_size;
    kupl_shm_comm_h kupl_comm = NULL;
    kupl_shm_win_h win = NULL;
    int *local_buffer = NULL;

    // 4. 初始化MPI环境（必须放在最前面，MPI进程创建的基础）
    ret = MPI_Init(&argc, &argv);
    if (ret != MPI_SUCCESS) {
        fprintf(stderr, "Error: MPI_Init failed\n");
        return -1;
    }

    // 5. 获取当前进程的rank和进程组大小（进程数）
    MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
    MPI_Comm_size(MPI_COMM_WORLD, &world_size);

    // 检查进程数是否为3（本次实验要求3个MPI进程）
    if (world_size != 3) {
        if (world_rank == 0) {
            fprintf(stderr, "Error: This experiment requires 3 MPI processes, but got %d\n", world_size);
        }
        MPI_Abort(MPI_COMM_WORLD, -1); // 终止所有MPI进程
        return -1;
    }

    // 6. 配置KUPL回调函数结构体
    kupl_shm_oob_cb_t oob_cbs;
    kupl_shm_oob_cb_h oob_cbs_h = &oob_cbs;
    oob_cbs_h->oob_allgather = oob_allgather_callback; // 绑定地址收集回调
    oob_cbs_h->oob_barrier = oob_barrier_callback;     // 绑定进程同步回调

    // 7. 创建KUPL通信域（核心步骤）
    ret = kupl_shm_comm_create(world_size, world_rank, getpid(), oob_cbs_h, (void *)MPI_COMM_WORLD, &kupl_comm);
    if (ret != 0 || kupl_comm == NULL) {
        fprintf(stderr, "Error: kupl_shm_comm_create failed, ret = %d\n", ret);
        MPI_Abort(MPI_COMM_WORLD, -1);
        return -1;
    }
    if (world_rank == 0) {
        printf("KUPL comm created successfully, world_size = %d\n", world_size);
    }

    // 8. 申请共享内存，创建内存窗口（测试小、中、大三种消息大小）
    // 这里以小消息为例，可替换为BUF_SIZE_MEDIUM或BUF_SIZE_LARGE测试不同场景
    int buf_size = BUF_SIZE_SMALL;
    ret = kupl_shm_win_alloc(buf_size * sizeof(int), kupl_comm, (void *)&local_buffer, &win);
    if (ret != 0 || win == NULL || local_buffer == NULL) {
        fprintf(stderr, "Error: kupl_shm_win_alloc failed, ret = %d\n", ret);
        kupl_shm_comm_destroy(kupl_comm); // 销毁已创建的通信域
        MPI_Abort(MPI_COMM_WORLD, -1);
        return -1;
    }

    // 9. 进程间数据拷贝（核心功能）
    // 9.1 rank0进程：初始化本地缓冲区，写入测试数据
    if (world_rank == 0) {
        for (int i = 0; i < buf_size; i++) {
            local_buffer[i] = i + 100; // 写入数据：100, 101, 102, 103（小消息场景）
        }
        printf("Rank 0: Data initialized successfully, buf_size = %d\n", buf_size);
        printf("Rank 0: Send data: ");
        for (int i = 0; i < buf_size; i++) {
            printf("%d ", local_buffer[i]);
        }
        printf("\n");

        // 9.2 rank0查询rank1和rank2的共享内存地址，拷贝数据
        void *remote_buf1 = NULL; // rank1的共享内存地址
        void *remote_buf2 = NULL; // rank2的共享内存地址

        // 查询rank1的共享内存地址
        ret = kupl_shm_win_query(win, 1, &remote_buf1);
        if (ret != 0 || remote_buf1 == NULL) {
            fprintf(stderr, "Error: kupl_shm_win_query failed for rank 1\n");
            goto cleanup; // 跳转到资源销毁步骤
        }

        // 查询rank2的共享内存地址
        ret = kupl_shm_win_query(win, 2, &remote_buf2);
        if (ret != 0 || remote_buf2 == NULL) {
            fprintf(stderr, "Error: kupl_shm_win_query failed for rank 2\n");
            goto cleanup;
        }

        // 数据拷贝（直接调用memcpy，无需MPI接口）
        memcpy(remote_buf1, local_buffer, buf_size * sizeof(int));
        memcpy(remote_buf2, local_buffer, buf_size * sizeof(int));
        printf("Rank 0: Data copied to rank 1 and rank 2 successfully\n");
    }

    // 10. 进程同步（确保rank1和rank2在读取数据前，rank0已完成拷贝）
    ret = oob_barrier_callback((void *)MPI_COMM_WORLD);
    if (ret != 0) {
        fprintf(stderr, "Error: oob_barrier_callback failed in main\n");
        goto cleanup;
    }

    // 11. rank1和rank2进程：读取共享内存数据，并进行校验
    if (world_rank == 1 || world_rank == 2) {
        printf("Rank %d: Received data: ", world_rank);
        for (int i = 0; i < buf_size; i++) {
            printf("%d ", local_buffer[i]);
        }
        printf("\n");

        // 数据校验（与预期数据对比，预期数据为100~103）
        int expected_buf[BUF_SIZE_SMALL] = {100, 101, 102, 103};
        ret = data_verify(expected_buf, local_buffer, buf_size);
        if (ret == 0) {
            printf("Rank %d: Data verify success\n", world_rank);
        } else {
            printf("Rank %d: Data verify failed\n", world_rank);
        }
    }

    // 12. 资源销毁（按顺序：先释放内存窗口，再销毁通信域，最后终止MPI）
cleanup:
    if (win != NULL) {
        kupl_shm_win_free(win); // 释放共享内存窗口
        win = NULL;
    }
    if (kupl_comm != NULL) {
        kupl_shm_comm_destroy(kupl_comm); // 销毁KUPL通信域
        kupl_comm = NULL;
    }
    local_buffer = NULL; // 释放本地缓冲区指针

    // 终止MPI环境
    MPI_Finalize();
    if (world_rank == 0) {
        printf("Experiment completed successfully, all resources released\n");
    }
    return 0;
}

#include <stdio.h>
#include <mpi.h>
#include "kupl_shm.h"

// oob_allgather回调：收集所有进程的内存地址（补充错误处理和多数据类型支持）
static int oob_allgather_callback(const void *sendbuf, void *recvbuf, int size, void *group, kupl_shm_datatype_t datatype) {
    int ret;
    // 1. 数据类型匹配：KUPL数据类型与MPI数据类型对应，确保数据拷贝正确
    switch (datatype) {
        case KUPL_SHM_DATATYPE_CHAR:
            // MPI_Allgather：将每个进程的sendbuf数据收集到所有进程的recvbuf
            ret = MPI_Allgather(sendbuf, size, MPI_CHAR, recvbuf, size, MPI_CHAR, (MPI_Comm)group);
            break;
        case KUPL_SHM_DATATYPE_INT:
            ret = MPI_Allgather(sendbuf, size, MPI_INT, recvbuf, size, MPI_INT, (MPI_Comm)group);
            break;
        case KUPL_SHM_DATATYPE_DOUBLE:
            ret = MPI_Allgather(sendbuf, size, MPI_DOUBLE, recvbuf, size, MPI_DOUBLE, (MPI_Comm)group);
            break;
        default:
            fprintf(stderr, "Error: Unsupported datatype in oob_allgather_callback\n");
            return -1; // 不支持的数据类型，返回失败
    }
    // 2. 错误检查：判断MPI接口调用是否成功，避免后续流程异常
    if (ret != MPI_SUCCESS) {
        fprintf(stderr, "Error: MPI_Allgather failed in oob_allgather_callback\n");
        return -1;
    }
    return 0; // 执行成功，返回0
}

// oob_barrier回调：实现进程同步，避免数据竞争
static int oob_barrier_callback(void *group) {
    // MPI_Barrier：阻塞当前进程，直到组内所有进程都调用该接口
    int ret = MPI_Barrier((MPI_Comm)group);
    if (ret != MPI_SUCCESS) {
        fprintf(stderr, "Error: MPI_Barrier failed in oob_barrier_callback\n");
        return -1;
    }
    return 0;
}

#include <stdio.h>
#include <stdlib.h>
#include <mpi.h>
#include "kupl_shm.h"

int main(int argc, char **argv) {
    int ret;
    int world_rank, world_size;
    kupl_shm_comm_h kupl_comm = NULL; // KUPL通信域句柄，初始化为NULL
    kupl_shm_oob_cb_t oob_cbs;
    kupl_shm_oob_cb_h oob_cbs_h = &oob_cbs;

    // 1. 初始化MPI环境，获取进程rank和进程数
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
    MPI_Comm_size(MPI_COMM_WORLD, &world_size);

    // 2. 绑定回调函数（必须在创建通信域前完成）
    oob_cbs_h->oob_allgather = oob_allgather_callback;
    oob_cbs_h->oob_barrier = oob_barrier_callback;

    // 3. 创建KUPL通信域（核心接口：kupl_shm_comm_create）
    ret = kupl_shm_comm_create(
        world_size,        // 参数1：通信组内的进程数（与MPI进程数一致）
        world_rank,        // 参数2：当前进程在通信组内的rank（与MPI rank一致）
        getpid(),          // 参数3：当前进程的PID（用于标识进程）
        oob_cbs_h,         // 参数4：回调函数结构体句柄
        (void *)MPI_COMM_WORLD, // 参数5：MPI通信域（与回调函数中的group对应）
        &kupl_comm         // 参数6：输出参数，存储创建的KUPL通信域句柄
    );

    // 4. 错误检查：判断通信域是否创建成功
    if (ret != 0 || kupl_comm == NULL) {
        fprintf(stderr, "Error: kupl_shm_comm_create failed, ret = %d\n", ret);
        MPI_Abort(MPI_COMM_WORLD, -1); // 终止所有MPI进程，避免资源泄漏
        return -1;
    }

    // 后续操作...

    // 销毁通信域