/*
 * Copyright (c) 2016-2020, NVIDIA CORPORATION. All rights reserved.
 *
 * See COPYRIGHT for license information
 */

#include "nvshmem.h"
#include "nvshmemx_error.h"
#include "nvshmem_internal.h"
#include "transport.h"
#include "util.h"
#include "topo.h"
#include <list>

/* Enumeration of possible PCIe paths and sister arrays for perf characteristics and string representations */
enum pci_distance { PATH_PIX = 0, PATH_PXB = 1, PATH_PHB = 2, PATH_NODE = 3, PATH_SYS = 4, PATH_COUNT = 5 };
static const int pci_distance_perf[PATH_COUNT] = { 4, 4, 3, 2, 1 };
static const char *pci_distance_string[PATH_COUNT] = { "PIX", "PXB", "PHB", "NODE", "SYS" };

bool nvshmemi_is_mpg_run = 0;

int nvshmemi_build_transport_map(nvshmemi_state_t *state) {
    int status = 0;
    int *local_map = NULL;

    state->transport_map = (int *)calloc(state->npes * state->npes, sizeof(int));
    NULL_ERROR_JMP(state->transport_map, status, NVSHMEMX_ERROR_OUT_OF_MEMORY, out,
                   "access map allocation failed \n");

    local_map = (int *)calloc(state->npes, sizeof(int));
    NULL_ERROR_JMP(local_map, status, NVSHMEMX_ERROR_OUT_OF_MEMORY, out,
                   "access map allocation failed \n");

    state->transport_bitmap = 0;

    for (int i = 0; i < state->npes; i++) {
        int reach_any = 0;

        for (int j = 0; j < NVSHMEM_TRANSPORT_COUNT; j++) {
            int reach = 0;

            if (!state->transports[j]) {
                continue;
            }

            status = state->transports[j]->host_ops.can_reach_peer(&reach, &state->pe_info[i],
                                                                   state->transports[j]);
            NZ_ERROR_JMP(status, NVSHMEMX_ERROR_INTERNAL, out, "can reach peer failed \n");
            INFO(NVSHMEM_TOPO, "[%d] reach %d to peer %d over transport %d", state->mype, reach,
                 i, j);

            state->transports[j]->cap[i] = reach;
            reach_any |= reach;

            if (reach) {
                int m = 1 << j;
                local_map[i] |= m;
                /* Add transport to the bitmap if this is the first PE to use it. */
                if ((state->transport_bitmap & m) == 0) {
                    state->transport_bitmap |= m;
                }
            }
        }

        if ((!reach_any) && (!nvshmemi_options.BYPASS_ACCESSIBILITY_CHECK)) {
            status = NVSHMEMX_ERROR_NOT_SUPPORTED;
            fprintf(stderr, "%s:%d: [GPU %d] Peer GPU %d is not accessible, exiting ... \n",
                    __FILE__, __LINE__, state->mype, i);
            goto out;
        }
    }
    INFO(NVSHMEM_TOPO, "[%d] transport bitmap: %x", state->mype, state->transport_bitmap);

    status = nvshmemi_boot_handle.allgather((void *)local_map, (void *)state->transport_map,
                                          sizeof(int) * state->npes, &nvshmemi_boot_handle);
    NZ_ERROR_JMP(status, NVSHMEMX_ERROR_INTERNAL, out, "allgather of ipc handles failed \n");

out:
    if (local_map) free(local_map);
    if (status) {
        if (state->transport_map) free(state->transport_map);
    }
    return status;
}

int nvshmemi_get_pcie_attrs(pcie_id_t *pcie_id, CUdevice cudev) {
    int status = 0;

    status = cuDeviceGetAttribute(&pcie_id->dev_id, CU_DEVICE_ATTRIBUTE_PCI_DEVICE_ID, cudev);
    NE_ERROR_JMP(status, CUDA_SUCCESS, NVSHMEMX_ERROR_INTERNAL, out,
                 "cuDeviceGetAttribute failed \n");

    status = cuDeviceGetAttribute(&pcie_id->bus_id, CU_DEVICE_ATTRIBUTE_PCI_BUS_ID, cudev);
    NE_ERROR_JMP(status, CUDA_SUCCESS, NVSHMEMX_ERROR_INTERNAL, out,
                 "cuDeviceGetAttribute failed \n");

    status = cuDeviceGetAttribute(&pcie_id->domain_id, CU_DEVICE_ATTRIBUTE_PCI_DOMAIN_ID, cudev);
    NE_ERROR_JMP(status, CUDA_SUCCESS, NVSHMEMX_ERROR_INTERNAL, out,
                 "cuDeviceGetAttribute failed \n");

out:
    return status;
}

int nvshmemi_detect_same_device(nvshmemi_state_t *state) {
    int status = 0;
    nvshmem_transport_pe_info_t my_info;

    my_info.pe = state->mype;
    status = nvshmemi_get_pcie_attrs(&my_info.pcie_id, state->cudevice);
    NE_ERROR_JMP(status, CUDA_SUCCESS, NVSHMEMX_ERROR_INTERNAL, out, "getPcieAttrs failed \n");

    my_info.hostHash = getHostHash();
    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, state->device_id);
    my_info.gpu_uuid = prop.uuid;

    //TODO: move this to a topo init function as it is reused in other functions in topo that follow
    state->pe_info =
        (nvshmem_transport_pe_info_t *)malloc(sizeof(nvshmem_transport_pe_info_t) * state->npes);
    NULL_ERROR_JMP(state->pe_info, status, NVSHMEMX_ERROR_OUT_OF_MEMORY, out,
                   "topo init info allocation failed \n");

    status = nvshmemi_boot_handle.allgather((void *)&my_info, (void *)state->pe_info,
                                          sizeof(nvshmem_transport_pe_info_t), &nvshmemi_boot_handle);
    NZ_ERROR_JMP(status, NVSHMEMX_ERROR_INTERNAL, out, "allgather of ipc handles failed \n");

    for (int i = 0; i < state->npes; i++) {
        (state->pe_info + i)->pe = i;
        if (i == state->mype) continue;

        status = (((state->pe_info + i)->hostHash == my_info.hostHash) &&
                  ((state->pe_info + i)->pcie_id.dev_id == my_info.pcie_id.dev_id) &&
                  ((state->pe_info + i)->pcie_id.bus_id == my_info.pcie_id.bus_id) &&
                  ((state->pe_info + i)->pcie_id.domain_id == my_info.pcie_id.domain_id));
        if (status) {
            INFO(NVSHMEM_INIT, "More than 1 PE per GPU detected. This is an MPG run.\n");
#if defined(NVSHMEM_PPC64LE)
            ERROR_EXIT("MPG support is currently not available on P9 platforms");
#endif
            nvshmemi_is_mpg_run = 1;
            status = 0;
        }
        /*NZ_ERROR_JMP(status, NVSHMEMX_ERROR_NOT_SUPPORTED, out,
                     "two PEs sharing a GPU is not supported \n");*/
    }

out:
    if (status) {
        state->cucontext = NULL;
        if (!state->pe_info) free(state->pe_info);
    }
    return status;
}

static int get_cuda_bus_id(int cuda_dev, char *bus_id) {
    int status = NVSHMEMX_SUCCESS;
    CUresult cu_err;

    cu_err = cuDeviceGetPCIBusId(bus_id, MAX_BUSID_SIZE, cuda_dev);
    if (cu_err != CUDA_SUCCESS) {
        ERROR_PRINT("cuDeviceGetPCIBusId failed with error: %d \n", cu_err);
        status = NVSHMEMX_ERROR_INTERNAL;
        goto out;
    }

out:
    return status;
}

static int get_device_path(char *bus_id, char **path) {
    int status = NVSHMEMX_SUCCESS;
    char pathname[MAXPATHSIZE + 1];
    char *cuda_rpath;
    char bus_path[] = "/sys/class/pci_bus/0000:00/device";

    for (int i = 0; i < 16; i++) bus_id[i] = tolower(bus_id[i]);
    memcpy(bus_path + sizeof("/sys/class/pci_bus/") - 1, bus_id, sizeof("0000:00") - 1);

    cuda_rpath = realpath(bus_path, NULL);
    NULL_ERROR_JMP(cuda_rpath, status, NVSHMEMX_ERROR_OUT_OF_MEMORY, out, "realpath failed \n");

    strncpy(pathname, cuda_rpath, MAXPATHSIZE);
    strncpy(pathname + strlen(pathname), "/", MAXPATHSIZE - strlen(pathname));
    strncpy(pathname + strlen(pathname), bus_id, MAXPATHSIZE - strlen(pathname));
    free(cuda_rpath);

    *path = realpath(pathname, NULL);
    NULL_ERROR_JMP(*path, status, NVSHMEMX_ERROR_OUT_OF_MEMORY, out, "realpath failed \n");

out:
    return status;
}

static enum pci_distance get_pci_distance(char *cuda_path, char *mlx_path) {
    int score = 0;
    int depth = 0;
    int same = 1;
    size_t i;
    for (i = 0; i < strlen(cuda_path); i++) {
        if (cuda_path[i] != mlx_path[i]) same = 0;
        if (cuda_path[i] == '/') {
            depth++;
            if (same == 1) score++;
        }
    }
    if (score <= 3) {
        /* Split the former PATH_SOC distance into PATH_NODE and PATH_SYS based on numaId */
        int numaId1 = getNumaId(cuda_path);
        int numaId2 = getNumaId(mlx_path);
        return ((numaId1 == numaId2) ? PATH_NODE : PATH_SYS);
    }
    if (score == 4) return PATH_PHB;
    if (score == depth - 1) return PATH_PIX;
    return PATH_PXB;
}

typedef struct nvshmemi_path_pair_info {
    int pe_idx;
    int dev_idx;
    enum pci_distance pcie_distance;
} nvshmemi_path_pair_info_t;

int get_device_by_distance(int *device, nvshmemi_state_t *state, struct nvshmem_transport *tcurr) {
    struct dev_info {
        char *dev_path; 
        int use_count;
    } *dev_info_all = NULL;

    struct gpu_info {
        char gpu_bus_id[MAX_BUSID_SIZE];
    } gpu_info, *gpu_info_all = NULL;

    std::list<nvshmemi_path_pair_info_t> pe_dev_pairs;
    std::list<nvshmemi_path_pair_info_t>::iterator pairs_iter;

    char **cuda_device_paths = NULL;
    int *pe_selected_devices = NULL;
    enum pci_distance *pe_device_distance = NULL;
    int *used_devs = NULL;

    int ndev;
    int mype_index = -1, mydev_index = -1;
    int i, dev_id, pe_id;
    int devices_assigned = 0;
    int status = NVSHMEMX_ERROR_INTERNAL;

    status = tcurr->host_ops.get_device_count(&ndev, tcurr);
    NZ_ERROR_JMP(status, NVSHMEMX_ERROR_INTERNAL, out,
                 "transport devices (setup_connections) failed \n");

    /* Allocate data structures start */
    dev_info_all = (struct dev_info *)calloc(ndev , sizeof(struct dev_info));
    NULL_ERROR_JMP(dev_info_all, status, NVSHMEMX_ERROR_OUT_OF_MEMORY, out,
                   "dev_info_all allocation failed \n");

    //TODO: we only need to exchange info among intra-node ranks, 
    gpu_info_all = (struct gpu_info *)calloc(state->npes, sizeof(struct gpu_info));
    NULL_ERROR_JMP(gpu_info_all, status, NVSHMEMX_ERROR_OUT_OF_MEMORY, out,
                   "gpu_info_all allocation failed \n");

    cuda_device_paths = (char **)calloc(state->npes_node, sizeof(char *));
    NULL_ERROR_JMP(cuda_device_paths, status, NVSHMEMX_ERROR_OUT_OF_MEMORY, out,
                   "Unable to allocate memory for PE/NIC Mapping.\n");
    
    pe_selected_devices = (int *)calloc(state->npes_node, sizeof(int));
    NULL_ERROR_JMP(pe_selected_devices, status, NVSHMEMX_ERROR_OUT_OF_MEMORY, out,
                   "Unable to allocate memory for PE/NIC Mapping.\n");
    for(pe_id = 0; pe_id < state->npes_node; pe_id++) {
        pe_selected_devices[pe_id] = -1;
    }

    pe_device_distance = (enum pci_distance *)calloc(state->npes_node, sizeof(enum pci_distance));
    NULL_ERROR_JMP(pe_device_distance, status, NVSHMEMX_ERROR_OUT_OF_MEMORY, out,
                   "Unable to allocate memory for PE/NIC Mapping.\n");

    used_devs = (int *)calloc(ndev, sizeof(int));
    NULL_ERROR_JMP(used_devs, status, NVSHMEMX_ERROR_OUT_OF_MEMORY, out,
                   "Unable to allocate memory for PE/NIC Mapping.\n");
    /* Allocate data structures end */

    /* Gather GPU and NIC paths start */
    status = get_cuda_bus_id(state->cudevice, gpu_info.gpu_bus_id);
    NZ_ERROR_JMP(status, NVSHMEMX_ERROR_INTERNAL, out, "get cuda busid failed \n");

    status = nvshmemi_boot_handle.allgather((void *)&gpu_info, (void *)gpu_info_all,
                                          sizeof(struct gpu_info), &nvshmemi_boot_handle);
    NZ_ERROR_JMP(status, NVSHMEMX_ERROR_INTERNAL, out, "allgather of gpu_info failed \n");

    pe_id = 0;
    for (i = 0; i < state->npes; i++) {
        if (state->pe_info[i].hostHash != state->pe_info[state->mype].hostHash) continue;
        
        status = get_device_path(gpu_info_all[i].gpu_bus_id, &cuda_device_paths[pe_id]);
        NZ_ERROR_JMP(status, NVSHMEMX_ERROR_INTERNAL, out, "get cuda path failed \n");
        /* to get back to our PE after the algorithm finishes. */
        if (i == state->mype) {
            mype_index = pe_id;
        }

        pe_id++;
        if (pe_id == state->npes_node) {
            break;
        }
    }

    if (pe_id != state->npes_node || mype_index == -1) {
        ERROR_JMP(status, NVSHMEMX_ERROR_INTERNAL, out,
                  "Number of PEs found doesn't match the PE node count.\n");
    }

    for (i = 0; i < ndev; i++) {
        status = tcurr->host_ops.get_pci_path(i, &dev_info_all[i].dev_path, tcurr);
        NZ_ERROR_JMP(status, NVSHMEMX_ERROR_INTERNAL, out, "get device path failed \n");
    }
    /* Gather GPU and NIC paths end */

    /* Get path distances start */
    for (pe_id = 0; pe_id < state->npes_node; pe_id++) {
        for (dev_id = 0; dev_id < ndev; dev_id++) {
            enum pci_distance distance_compare;
            distance_compare = get_pci_distance(cuda_device_paths[pe_id], dev_info_all[dev_id].dev_path);
            if (unlikely(pe_dev_pairs.empty())) {
                pe_dev_pairs.push_front({pe_id, dev_id, distance_compare});
            } else {
                for(pairs_iter = pe_dev_pairs.begin(); pairs_iter != pe_dev_pairs.end(); pairs_iter++) {
                    if (distance_compare < (*pairs_iter).pcie_distance) {
                        break;
                    }
                }
                INFO(NVSHMEM_TOPO, "PE %d: %s dev %d: %s distance: %d\n",
                     pe_id, cuda_device_paths[pe_id],
                     dev_id, dev_info_all[dev_id].dev_path,
                     distance_compare);
                pe_dev_pairs.insert(pairs_iter, {pe_id, dev_id, distance_compare});
            }
        }
    }
    /* Get path distances end */

    /* loop one, do initial assignments of NICs */
    for (pairs_iter = pe_dev_pairs.begin(); pairs_iter != pe_dev_pairs.end(); pairs_iter++) {
        /* skip pairs where the GPU already has a partner in the first loop */
        if (pe_selected_devices[(*pairs_iter).pe_idx] != -1) {
            continue;
        }

        INFO(NVSHMEM_TOPO, "Pairing PE %d with device %d at distance %d\n",
             (*pairs_iter).pe_idx, (*pairs_iter).dev_idx, (*pairs_iter).pcie_distance);
        pe_selected_devices[(*pairs_iter).pe_idx] = (*pairs_iter).dev_idx;
        pe_device_distance[(*pairs_iter).pe_idx] = (*pairs_iter).pcie_distance;
        used_devs[(*pairs_iter).dev_idx]++;
        devices_assigned++;

        if (devices_assigned == state->npes_node) {
            break;
        }
    }

    /* loop two, load balance the NICs. */
    for (pe_id = 0; pe_id < state->npes_node; pe_id++) {
        int nic_density = used_devs[pe_selected_devices[pe_id]];

        /* Can't find a less populated NIC if ours is only assigned to 1 gpu. */
        if (nic_density < 2) {
            continue;
        }

        for (pairs_iter = pe_dev_pairs.begin(); pairs_iter != pe_dev_pairs.end(); pairs_iter++) {
            /* Never change for a less optimal NIC. */

            if ((*pairs_iter).pe_idx != pe_id) {
                continue;
            }

            if (pci_distance_perf[(*pairs_iter).pcie_distance] < pci_distance_perf[pe_device_distance[pe_id]]) {
                break;
            }

            if ((nic_density - used_devs[(*pairs_iter).dev_idx]) >= 2) {
                INFO(NVSHMEM_TOPO, "Re-Pairing PE %d with device %d at distance %d\n",
                    (*pairs_iter).pe_idx, (*pairs_iter).dev_idx, (*pairs_iter).pcie_distance);
                used_devs[pe_selected_devices[pe_id]]--;
                used_devs[(*pairs_iter).dev_idx]++;
                nic_density = used_devs[(*pairs_iter).dev_idx];
                pe_selected_devices[pe_id] = (*pairs_iter).dev_idx;
                pe_device_distance[pe_id] = (*pairs_iter).pcie_distance;
                if (nic_density < 2) {
                    break;
                }
            }
        }
    }

    mydev_index = pe_selected_devices[mype_index];
    if (used_devs[mydev_index] > 1) {
        INFO(NVSHMEM_TOPO, "Our PE is sharing its NIC with %d other PEs.\n", used_devs[mydev_index]);
    }

    if (pci_distance_perf[pe_device_distance[mype_index]] < pci_distance_perf[PATH_PIX]) {
        INFO(NVSHMEM_TOPO, "Our PE is connected to a NIC with pci distance %s." \
                            "this will provide less than optimal performance.\n",
                            pci_distance_string[pe_device_distance[mype_index]]);
    }

out:
    if (dev_info_all) {
        free(dev_info_all);
    }

    if (gpu_info_all) {
        free(gpu_info_all);
    }

    if (cuda_device_paths) {
        for (i = 0; i < state->npes_node; i++) {
            if (cuda_device_paths[i]) {
                free(cuda_device_paths[i]);
            }
        }
        free(cuda_device_paths);
    }

    pe_dev_pairs.clear();

    if (pe_selected_devices) {
        free(pe_selected_devices);
    }
    
    if (used_devs) {
        free(used_devs);
    }

    if (pe_device_distance) {
        free(pe_device_distance);
    }

    if (mydev_index >= 0) {
        *device = mydev_index;
        return 0;
    }

    return status;
}
