/*
 *  Copyright (C) 2026, Northwestern University
 *  See COPYRIGHT notice in top-level directory.
 */

#ifdef HAVE_CONFIG_H
# include <config.h>
#endif

#include <stdio.h>
#include <stdlib.h>
#include <string.h>     /* strlen(), strcat(), strerror() */
#include <fcntl.h>      /* open(), O_CREAT */
#include <sys/types.h>  /* open(), umask(), fstat() */

#if defined(HAVE_SYS_STAT_H) && HAVE_SYS_STAT_H == 1
#include <sys/stat.h>   /* fstat() */
#endif
#include <unistd.h>     /* fstat() */

#include <assert.h>
#include <sys/errno.h>

#include "gioi.h"

#define SET_INFO(fh) {                                                      \
    char int_str[16];                                                       \
                                                                            \
    /* set file hints with values only available after file is opened */    \
    MPI_Info_set(fh->info, "file_system_type", "UFS:");                     \
                                                                            \
    snprintf(int_str, 16, "%d", fh->hints->striping_unit);                  \
    MPI_Info_set(fh->info, "striping_unit", int_str);                       \
                                                                            \
    /* collective buffer size must be at least file striping size */        \
    if (fh->hints->cb_buffer_size < fh->hints->striping_unit) {             \
        fh->hints->cb_buffer_size = fh->hints->striping_unit;               \
        snprintf(int_str, 16, " %d", fh->hints->cb_buffer_size);            \
        MPI_Info_set(fh->info, "cb_buffer_size", int_str);                  \
    }                                                                       \
}

#if GIO_MPI_ENABLED == 1

/*----< UFS_set_cb_node_list() >---------------------------------------------*/
/* Construct the list of I/O aggregators. It sets the followings.
 *   fh->hints->aggr_ranks[].
 *   fh->hints->cb_nodes and set file info for hint cb_nodes.
 *   fh->is_agg: indicating whether this rank is an I/O aggregator
 *   fh->my_cb_nodes_index: index into fh->hints->aggr_ranks[]. -1 if N/A
 */
static
int UFS_set_cb_node_list(GIO_File fh)
{
    char value[MPI_MAX_INFO_VAL + 1], int_str[16];
    int i, j, k, nprocs, rank, *nprocs_per_node, **ranks_per_node;

    MPI_Comm_size(fh->comm, &nprocs);
    MPI_Comm_rank(fh->comm, &rank);

    if (fh->hints->cb_nodes == 0)
        /* If hint cb_nodes is not set by user, select one rank per node to be
         * an I/O aggregator
         */
        fh->hints->cb_nodes = fh->num_NUMAs;
    else if (fh->hints->cb_nodes > nprocs)
        /* cb_nodes must be <= nprocs */
        fh->hints->cb_nodes = nprocs;

    fh->hints->aggr_ranks = (int*) GIOI_Malloc(sizeof(int) * fh->hints->cb_nodes);
    if (fh->hints->aggr_ranks == NULL)
        return GIO_ENOMEM;

    /* number of MPI processes running on each NUMA node */
    nprocs_per_node = (int*) GIOI_Calloc(fh->num_NUMAs, sizeof(int));

    for (i=0; i<nprocs; i++) nprocs_per_node[fh->NUMA_IDs[i]]++;

    /* construct rank IDs of MPI processes running on each node */
    ranks_per_node = (int**) GIOI_Malloc(sizeof(int*) * fh->num_NUMAs);
    ranks_per_node[0] = (int*) GIOI_Malloc(sizeof(int) * nprocs);
    for (i=1; i<fh->num_NUMAs; i++)
        ranks_per_node[i] = ranks_per_node[i - 1] + nprocs_per_node[i - 1];

    for (i=0; i<fh->num_NUMAs; i++) nprocs_per_node[i] = 0;

    /* Populate ranks_per_node[], list of MPI ranks running on each NUMA node.
     * Populate nprocs_per_node[], number of MPI processes on each NUMA node.
     */
    for (i=0; i<nprocs; i++) {
        k = fh->NUMA_IDs[i];
        ranks_per_node[k][nprocs_per_node[k]] = i;
        nprocs_per_node[k]++;
    }

#ifdef ROUND_ROBIN_POLICY
    /* Round-robin assignment policy selects aggregators in the round-robin
     * fashion across compute nodes, i.e. consecutive aggregators are selected
     * from different, consecutive nodes.
     */
    k = j = 0;
    for (i=0; i<fh->hints->cb_nodes; i++) {
        if (j >= nprocs_per_node[k]) { /* if run out of ranks in this node k */
            k++;
            if (k == fh->num_NUMAs) { /* round-robin to first node */
                k = 0;
                j++;
            }
        }
        /* select the jth rank of node k as an I/O aggregator */
        fh->hints->aggr_ranks[i] = ranks_per_node[k++][j];
        if (rank == fh->hints->aggr_ranks[i]) {
            fh->is_agg = 1;
            fh->my_cb_nodes_index = i;
        }
        if (k == fh->num_NUMAs) { /* round-robin to first node */
            k = 0;
            j++;
        }
    }
#else
    /* Block assignment policy selects aggregators from compute nodes in the
     * block fashion, i.e. consecutive aggregators are selected from the same
     * node.
     *
     * Performance evaluation on Perlmutter Lustre using WRF-IO does not show a
     * noticeable difference between the round-robin and block policies.
     */
    int avg = fh->hints->cb_nodes / fh->num_NUMAs;
    int rem = fh->hints->cb_nodes % fh->num_NUMAs;
    k = 0;
    for (i=0; i<fh->num_NUMAs; i++) {
        int num_aggr = (i < rem) ? avg + 1 : avg;
        /* pick num_aggr processes as I/O aggregators in this node i */
        for (j=0; j<num_aggr; j++) {
            fh->hints->aggr_ranks[k] = ranks_per_node[i][j];
            if (rank == fh->hints->aggr_ranks[k]) {
                fh->is_agg = 1;
                fh->my_cb_nodes_index = k;
            }
            k++;
        }
    }
#endif
    GIOI_Free(ranks_per_node[0]);
    GIOI_Free(ranks_per_node);
    GIOI_Free(nprocs_per_node);

    /* Set file hints with values only available after cb_nodes and
     * aggr_ranks[] have been established.
     */
    snprintf(int_str, 16, "%d", fh->hints->cb_nodes);
    MPI_Info_set(fh->info, "cb_nodes", int_str);

    /* add hint "cb_node_list", list of aggregators' rank IDs */
    snprintf(value, 16, "%d", fh->hints->aggr_ranks[0]);
    for (i=1; i<fh->hints->cb_nodes; i++) {
        snprintf(int_str, 16, " %d", fh->hints->aggr_ranks[i]);
        if (strlen(value) + strlen(int_str) >= MPI_MAX_INFO_VAL-5) {
            strcat(value, " ...");
            break;
        }
        strcat(value, int_str);
    }
    MPI_Info_set(fh->info, "cb_node_list", value);

    return GIO_NOERR;
}
#endif

/*----< GIOI_UFS_open() >----------------------------------------------------*/
/*   1. root creates/opens the file first
 *   2. root obtains file striping info, i.e. statbuf.st_blksize
 *   3. root broadcasts striping info
 *   4. non-root processes receive striping info from root
 *   5. non-root processes opens the file
 */
int
GIOI_UFS_open(GIO_File fh)
{
    char *mode_str = (fh->amode & O_CREAT) ? "creat" : "open";
    int err=GIO_NOERR, rank, perm, old_mask;
    int stripe_size, stripin_info[2] = {GIO_NOERR, 1048576};

    MPI_Comm_rank(fh->comm, &rank);

    stripe_size = 1048576; /* default to 1 MiB */

    old_mask = umask(022);
    umask(old_mask);
    perm = old_mask ^ GIOI_PERM;

    /* Root process creates/opens the file first, obtains statbuf.st_blksize,
     * broadcasts it, and followed by the rest of processes open the file.
     */
    if (rank > 0) goto err_out;

    /* Root first creates/opens the file and obtains st_blksize */
    fh->fd_sys = open(fh->filename, fh->amode, perm);
    if (fh->fd_sys == -1) {
        fprintf(stderr, "%s line %d: rank %d failed to %s file %s (%s)\n",
                __func__,__LINE__, rank, mode_str, fh->filename,
                strerror(errno));
        err = GIOI_error_posix(mode_str);
        goto err_out;
    }
    fh->is_open = 1;

    /* Only root obtains the striping information and bcast to all other
     * processes. For UFS, file striping is the file system block size.
     */
    stripe_size = 1048576; /* default to 1 MiB */
#if defined(HAVE_SYS_STAT_H) && HAVE_SYS_STAT_H == 1
    /* Get the underlying file system block size as file striping_unit */
    struct stat statbuf;
    err = fstat(fh->fd_sys, &statbuf);
    if (err >= 0)
        /* file system block size usually < MAX_INT */
        stripe_size = (int)statbuf.st_blksize;
#endif

err_out:
    stripin_info[0] = err;
    stripin_info[1] = stripe_size;

    MPI_Bcast(stripin_info, 2, MPI_INT, 0, fh->comm);

    fh->hints->striping_unit = stripin_info[1];

    if (stripin_info[0] != GIO_NOERR) {
        /* root failed to create/open the file */
        fprintf(stderr, "%s line %d: root failed to %s UFS file %s\n",
                __FILE__, __LINE__, mode_str, fh->filename);
        return err;
    }

    SET_INFO(fh)

    /* Construct cb_nodes rank list, which requires fh->num_NUMAs to be known
     * on all processes.
     */
    UFS_set_cb_node_list(fh);

    /* Now non-root I/O aggregators and only aggregators open the file. */
    if (rank > 0 && fh->is_agg) {
        fh->fd_sys = open(fh->filename, fh->amode, perm);
        if (fh->fd_sys == -1) {
            fprintf(stderr, "%s line %d: rank %d failed to open file %s (%s)\n",
                    __func__,__LINE__, rank, fh->filename, strerror(errno));
            err = GIOI_error_posix("open");
        }
        fh->is_open = 1;
    }

    return err;
}

/*----< GIOI_UFS_open_on_demand() >------------------------------------------*/
/* This subroutine is an independent call.
 *
 * This subroutine is called by the non-aggregators only. Its fh has been
 * allocated but fh->is_open is 0, i.e. the non-aggregator has not made the
 * system open() call to open the file. In this case, it calls open() and
 * retrieve file striping size.
 */
int GIOI_UFS_open_on_demand(GIO_File fh)
{
    int err=GIO_NOERR, rank, perm, old_mask;

#if GIO_DEBUG_MODE == 1
    assert(fh != NULL);
    assert(fh->is_open == 0);
#endif

    old_mask = umask(022);
    umask(old_mask);
    perm = old_mask ^ GIOI_PERM;

    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    /* open the file now */
    fh->fd_sys = open(fh->filename, fh->amode, perm);
    if (fh->fd_sys == -1) {
        fprintf(stderr,"%s line %d: world rank %d failed to open file %s (%s)\n",
                __FILE__,__LINE__, rank, fh->filename, strerror(errno));
        return GIOI_error_posix("open");
    }
    fh->is_open = 1;

    fh->hints->striping_unit = 1048576; /* default to 1 MiB */

#if defined(HAVE_SYS_STAT_H) && HAVE_SYS_STAT_H == 1
    /* Get the underlying file system block size as file striping_unit */
    struct stat statbuf;
    err = fstat(fh->fd_sys, &statbuf);
    if (err >= 0)
        /* file system block size usually < MAX_INT */
        fh->hints->striping_unit = (int)statbuf.st_blksize;
#endif

    return err;
}

