#include <stdio.h>
#include <stdlib.h>
#include <level_zero/ze_api.h>
#include <mpi.h>

#define ZE_CHECK(result)                                                   \
    if (result != ZE_RESULT_SUCCESS) {                                      \
        fprintf(stderr, "Error: %s:%d\n", __FILE__, __LINE__);              \
        exit(EXIT_FAILURE);                                                 \
    }

int main() {

    int rank, size;
    // Initialize MPI
    MPI_Init(NULL, NULL);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    ze_result_t result;
    ze_driver_handle_t driver;
    ze_context_handle_t context;
    ze_device_handle_t device[4];

    // Initialize Level Zero
    result = zeInit(ZE_INIT_FLAG_GPU_ONLY);
    ZE_CHECK(result);

    // Get the driver
    uint32_t driverCount = 0;
    result = zeDriverGet(&driverCount, NULL);
    ZE_CHECK(result);

    result = zeDriverGet(&driverCount, &driver);
    ZE_CHECK(result);

    // Create the context
    ze_context_desc_t contextDesc = {ZE_STRUCTURE_TYPE_CONTEXT_DESC, NULL, 0};
    result = zeContextCreate(driver, &contextDesc, &context);
    ZE_CHECK(result);

    // Get the device
    uint32_t deviceCount = 0;
    result = zeDeviceGet(driver, &deviceCount, NULL);
    ZE_CHECK(result);

    result = zeDeviceGet(driver, &deviceCount, device);
    ZE_CHECK(result);

    // Generate a memory allocation on the device
    const size_t bufferSize = 1024;
    void* deviceMemory = NULL;

    ze_device_mem_alloc_desc_t deviceDesc = {
        ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC,
        NULL,
        ZE_DEVICE_MEM_ALLOC_FLAG_BIAS_CACHED,
        0
    };
    ze_device_mem_alloc_desc_t otherDeviceDesc = {
        ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC,
        NULL,
        ZE_DEVICE_MEM_ALLOC_FLAG_BIAS_CACHED,
        0
    };

    ze_ipc_mem_handle_t ipcMemHandle;

    if (rank == 0) {
        ZE_CHECK(zeMemAllocDevice(context, &deviceDesc, bufferSize, 1, device[rank], &deviceMemory));

        // Get the IPC handle from the device memory
        ZE_CHECK(zeMemGetIpcHandle(context, deviceMemory, &ipcMemHandle));

        MPI_Send(&ipcMemHandle, sizeof(ze_ipc_mem_handle_t), MPI_BYTE, 1, 0, MPI_COMM_WORLD);
        MPI_Send(&context, sizeof(ze_context_handle_t), MPI_BYTE, 1, 0, MPI_COMM_WORLD);
    } else {
        MPI_Recv(&ipcMemHandle, sizeof(ze_ipc_mem_handle_t), MPI_BYTE, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

        // Open the IPC memory handle in another process
        ze_device_handle_t otherDevice;
        ze_context_handle_t otherContext;
        MPI_Recv(&otherContext, sizeof(ze_context_handle_t), MPI_BYTE, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

        void* otherDeviceMemory = NULL;

        ZE_CHECK(zeMemOpenIpcHandle(context, device[rank], ipcMemHandle, otherDeviceDesc.flags, &otherDeviceMemory));
    }

    // Copy data from deviceMemory to otherDeviceMemory
    /*result = zeCommandListAppendMemoryCopy(
        nullptr, otherDeviceMemory, deviceMemory, bufferSize, nullptr, 0, nullptr);
    ZE_CHECK(result);*/

    // Close the IPC memory handle
    /*result = zeMemCloseIpcHandle(context, otherDeviceMemory);
    ZE_CHECK(result);*/

    // Free the device memory
    result = zeMemFree(context, deviceMemory);
    ZE_CHECK(result);

    // Destroy the context
    result = zeContextDestroy(context);
    ZE_CHECK(result);

    MPI_Finalize();

    return 0;
}