#include <unistd.h>
#include <sys/types.h>
#include <string.h>
#include <sys/resource.h>
#include <iostream>
#include <iomanip>
#include <sstream>
#include <chrono>
#include <ctime>
#include <cstdint>

#include "bootstrap.h"

namespace sccl {
namespace hardware {
namespace topology {
namespace bootstrap {

#define MAX_THREADS (128)

struct bootstrapRootArgs {
    scclSocket_t* listenSock;
    uint64_t magic;
};

// 构造函数
scclBootstrap::scclBootstrap(struct scclRankInfo* rank_info, struct scclBootstrapComm* comm) {
    // 初始化线程池
    int thread_cnt = ::std::min(MAX_THREADS, rank_info->nRanks);
    pthread_pool   = new ThreadPool(thread_cnt);

    scclResult_t res;
    // 将 handle 结构体清零
    SCCLCHECKGOTO(scclCalloc(&handle, 1), res, failure);

    // 初始化bootstrap网络环境
    SCCLCHECKGOTO(bootstrapInit(rank_info, comm), res, failure);

    return;

failure:
    WARN("bootstrap not implemented yet");
    return;
}

scclBootstrap::~scclBootstrap() {
    if(handle) {
        free(handle);
    }
    if(bootstrap_net) {
        delete bootstrap_net;
    }

    if(pthread_pool) {
        delete pthread_pool;
    }
}

/**
 * 初始化bootstrap通信环境
 *
 * @param rank_info 包含rank和nRanks信息的唯一标识符
 * @param comm 需要初始化的bootstrap通信结构体
 * @return scclResult_t 返回操作结果，成功返回scclSuccess
 *
 * 该函数负责：
 * 1. 设置WarpSize
 * 2. 初始化unique_info信息
 * 3. 创建并初始化bootstrap socket句柄
 */
scclResult_t scclBootstrap::bootstrapInit(const struct scclRankInfo* rank_info, struct scclBootstrapComm* comm) {
    // 如果已经初始化，直接返回成功
    if(asm_ops::ld_acquire_sys_global(&initialized))
        return scclSuccess;

    // 加锁以确保初始化过程的线程安全
    pthread_mutex_lock(&initLock);
    // 如果尚未初始化，进行初始化操作
    if(!initialized) {
        // 初始化通信结构的WarpSize
        comm->WarpSize = warpSize;
        // 获取CPU亲和性
        sched_getaffinity(0, sizeof(cpu_set_t), &comm->cpuAffinity);

        // 分配并初始化scclUniqueInfo信息
        SCCLCHECK(scclCalloc(&(comm->unique_info), 1));

        // 获取环境变量SCCL_NET_NAME的值，如果不存在则默认使用"IB"
        const char* envNetName = getenv("SCCL_NET_NAME");
        char* netName          = (envNetName != NULL) ? strdup(envNetName) : strdup("IB");
        // 打印网络名称
        printf("netName=%s\n", netName);

        // 初始化网络和引导网络
        SCCLCHECK(net::scclNetInit(netName, comm->scclNet));
        // 释放分配的网络名称字符串
        free(netName);

        // 调用bootstrapNetInit初始化CPU硬件，环境检查
        SCCLCHECK(bootstrapBasicInit());
        // 初始化唯一信息结构体
        SCCLCHECK(bootstrapUniqueInfoInit(rank_info, comm->scclNet, comm->unique_info));

        // 初始化网络结构体，用于bootstrap阶段的socket通信
        bootstrap_net = new bootstrapNet(comm);

        // 设置当前rank的socket信息给handle
        SCCLCHECK(getRandomData(&handle->magic, sizeof(handle->magic)));
        memcpy(&handle->addr, &bootstrap_net->bootstrapNetIfAddr, sizeof(scclSocketAddress_t));

        // SCCLCHECK(getIpcSocketAddr(&handle->peerIpcAddr));
#if 0
        // char line[100];
        // sprintf(line, "pos 55: rank=%d", rank_info->rank);
        // SCCLCHECK(net::printSocketAddr(&handle->addr, line));
#endif

        bootstrapAllGather(comm->unique_info);

        // 设置初始化完成标志
        asm_ops::st_release_sys_global(&initialized, true);
    }
    // 解锁
    pthread_mutex_unlock(&initLock);

    return scclSuccess;
}

/**
 * @brief 执行基本的引导程序初始化
 *
 * 该函数负责初始化引导程序的基本组件，包括网络引导和系统环境检查。
 * 使用互斥锁确保线程安全，避免重复初始化。
 *
 * @note 如果NUMA自动平衡已启用，会发出警告提示可能影响性能。
 * @note 会检查内核版本信息并记录。
 *
 * @return scclResult_t 返回初始化结果，成功返回scclSuccess
 */
scclResult_t scclBootstrap::bootstrapBasicInit() {
    // 始终初始化引导网络
    SCCLCHECK(bootstrap_net->bootstrapNetInit());
    // SCCLCHECK(scclNetPluginInit());  // collnet使用

    char strValue[1024];
    // 检查NUMA自动平衡是否启用
    SCCLCHECK(scclTopoGetStrFromSys("/proc/sys/kernel", "numa_balancing", strValue));
    if(strcmp(strValue, "1") == 0)
        WARN("NUMA自动平衡已启用，这可能导致RCCL性能的不稳定性！通过\"sudo sysctl kernel.numa_balancing=0\"禁用");

    // 获取内核版本信息
    SCCLCHECK(scclTopoGetStrFromSys("/proc", "version", strValue));

    char *verStr, *state;
    verStr = strtok_r(strValue, " ", &state);
    INFO(SCCL_LOG_BOOTSTRAP, "内核版本: %s", verStr);

    for(int i = 0; i < 2; i++) {
        verStr = strtok_r(NULL, " ", &state);
        if(verStr == NULL)
            break;
    }

// TODO: 最终确定是否需要检查版本信息
#if 0
        // 检查是否为Cray系统
        if(strstr(verStr, "cray") == NULL) {
            // 获取BIOS版本信息
            SCCLCHECK(scclTopoGetStrFromSys("/sys/devices/virtual/dmi/id", "bios_version", strValue));

            if(strncmp("Hyper-V UEFI Release", strValue, 20) != 0) {
                FILE* file;
                // 读取内核命令行参数
                if((file = fopen("/proc/cmdline", "r")) != NULL) {
                    if(feof(file) == 0 && ferror(file) == 0) {
                        int len       = fread(strValue, 1, 1024, file);
                        strValue[len] = '\0';
                    }
                    fclose(file);
                }
                // 检查是否缺少"iommu=pt"参数
                if(strstr(strValue, "iommu=pt") == NULL)
                    WARN("内核命令行中缺少\"iommu=pt\"参数，这可能导致系统不稳定或挂起！");
            }
#ifndef HIP_UNCACHED_MEMORY
            // 检查环境变量"HSA_FORCE_FINE_GRAIN_PCIE"
            char* env = getenv("HSA_FORCE_FINE_GRAIN_PCIE");
            printf("HSA env=%s\n", env);
            if(env == NULL || strcmp(env, "1") != 0)
                WARN("环境变量中缺少\"HSA_FORCE_FINE_GRAIN_PCIE=1\"，这可能导致RCCL性能低下，系统不稳定或挂起！");
#endif
            float* ptr;
            // 尝试分配细粒度PCIe内存
            hipError_t err = hipExtMallocWithFlags((void**)&ptr, 128, hipDeviceMallocFinegrained);
            if(err != hipSuccess)
                hsaFineGrainFlag = false;
        }
#endif

    return scclSuccess;
}

/**
 * @brief 初始化唯一信息结构体
 *
 * 该函数用于初始化scclUniqueInfo结构体，包括设置rank信息、设备信息、主机和进程哈希值，
 * 以及硬件相关信息（GPU、CPU、RDMA、PCI等）。
 *
 * @param rank_info 输入参数，包含rank相关信息
 * @param unique_info 输出参数，待初始化的唯一信息结构体
 * @return scclResult_t 返回操作结果，scclSuccess表示成功
 */
scclResult_t scclBootstrap::bootstrapUniqueInfoInit(const struct scclRankInfo* rank_info, scclNet_t* scclNet, struct scclUniqueInfo* unique_info) {
    ////////////////// 设置id信息 //////////////////
    unique_info->rank       = rank_info->rank;       // 将unique_id的rank赋值给unique_info的rank
    unique_info->nRanks     = rank_info->nRanks;     // 将unique_id的nRanks赋值给unique_info的nRanks
    unique_info->localRanks = rank_info->localRanks; // 将unique_id的localRanks赋值给unique_info的localRanks
    unique_info->localRank  = rank_info->localRank;  // 计算unique_info的localRank

    LECHECK(unique_info->localRanks, unique_info->localRank); // 检查localRank是否小于localRanks

    uint32_t devices_num;
    SCCLCHECK(rocm_smi_init());                     // 初始化ROCM SMI库
    SCCLCHECK(rocm_smi_getNumDevice(&devices_num)); // 获取设备数量
    LTCHECK(devices_num, 0);                        // 检查设备数量是否大于0

    unique_info->deviceCnt = static_cast<int>(devices_num); // 将设备数量转换为int并赋值给unique_info的deviceCnt
    LECHECK(unique_info->deviceCnt, unique_info->hipDev);   // 检查hipDev是否小于deviceCnt
    HIPCHECK(hipSetDevice(unique_info->hipDev));            // 设置当前设备为hipDev
    unique_info->hipDev = rank_info->hipDev;

    // 获取其他基础信息
    unique_info->hostHash = getHostHash(); // 获取主机哈希值并赋值给unique_info的hostHash
    unique_info->pidHash  = getPidHash();  // 获取进程ID哈希值并赋值给unique_info的pidHash

    ////////////////// 设置硬件信息 //////////////////
    struct topoLocalNode* p_localNode = &unique_info->localNode;

    // 设置GPU信息
    p_localNode->gpu.dev = rank_info->hipDev;
    hipDeviceProp_t deviceProp;
    HIPCHECK(hipGetDeviceProperties(&deviceProp, rank_info->hipDev));
    snprintf(p_localNode->gpu.name, sizeof(p_localNode->gpu.name), "%s", deviceProp.name);
    snprintf(p_localNode->gpu.gcn, sizeof(p_localNode->gpu.gcn), "%s", deviceProp.gcnArchName);
    p_localNode->gpu.compCap = deviceProp.major * 10 + deviceProp.minor;

    // 设置CPU信息
    memcpy(&p_localNode->cpu.socketAddr, &bootstrap_net->bootstrapNetIfAddr, sizeof(scclSocketAddress_t));

    // 设置RDMA信息
    SCCLCHECK(scclNet->getProperties(rank_info->hipDev, &p_localNode->net.props));
    SCCLCHECK(scclNet->devices(&p_localNode->net.count));

    // 设置PCI信息
    SCCLCHECK(getBusId(rank_info->hipDev, &p_localNode->pci.busId));

#if 1
    printf("topoLocalNode size=%ld\n", sizeof(struct topoLocalNode));
    SCCLCHECK(net::printNetProps(&p_localNode->net.props, rank_info->rank, rank_info->localRank));
#endif

    return scclSuccess;
}

/**
 * 检查bootstrap是否已成功初始化
 *
 * 该函数通过检查handle指针和initialized标志来验证初始化状态
 * 使用互斥锁确保线程安全
 *
 * @return scclSuccess 如果已初始化成功
 * @return scclSystemError 如果未初始化或初始化失败
 */
scclResult_t scclBootstrap::bootstrapInitCheck() {
    scclResult_t res = scclSuccess;

    // 加锁以确保初始化过程的线程安全
    pthread_mutex_lock(&initLock);
    if(handle == nullptr || initialized == false) {
        res = scclSystemError;
    }
    pthread_mutex_unlock(&initLock); // 解锁

    return res;
}

scclResult_t scclBootstrap::bootstrapAllGather(struct scclUniqueInfo* unique_info) {
    // 1.节点内通信 allgather

    // 2.节点间通信，ring allgather

    // 3.节点内通信 allgather

    return scclSuccess;
}

/////////////////////////////////////////////////////////////////////////////////////////////
// // 将本地socket地址写入到/tmp/文件夹的文件中，通过nfs共享存储，其他rank可见
// scclResult_t bootstrapGetAllNodes(const struct scclUniqueInfo* unique_info, struct scclBootstrapComm* comm) {
//     // // 分配并初始化IPC套接字
//     // struct scclIpcSocket ipcSock = {0};
//     // // Create a UDS socket to receive the converted fd
//     // SCCLCHECK(scclIpcSocketInit(&ipcSock, unique_info->rank, /*hash*/ handle->magic, /*abortFlag*/ NULL));
//     // printf("fd=%d, socketName=%s\n", ipcSock.fd, ipcSock.socketName);
//     return scclInProgress;
// }

} // namespace bootstrap
} // namespace topology
} // namespace hardware
} // namespace sccl