#pragma once

#include <assert.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <poll.h>
#include <sys/types.h>
#include <unistd.h>

#include "ibvwrap.h"
#include "socket.h"
#include "net_utils.h"

namespace sccl {
namespace hardware {
namespace net {
namespace net_ib {

/*IB的通信状态*/
enum scclIbCommState : uint8_t {
    scclIbCommStateStart        = 0, // 初始状态
    scclIbCommStateConnect      = 1, // 尝试连接状态
    scclIbCommStateAccept       = 3, // 接受连接状态
    scclIbCommStateSend         = 4, // 发送数据状态
    scclIbCommStateRecv         = 5, // 接收数据状态
    scclIbCommStateConnecting   = 6, // 正在连接状态
    scclIbCommStateConnected    = 7, // 已连接状态
    scclIbCommStatePendingReady = 8, // 等待准备状态
};

/*通信的阶段*/
struct scclIbCommStage {
    enum scclIbCommState state; // 通信阶段的状态
    int offset;                 // 数据偏移量
    void* buffer;               // 用于通信的缓冲区指针
    void* comm;                 // 通信对象指针
};

/*监听通信的上下文*/
struct scclIbListenComm {
    int dev;                            // 设备标识符
    struct net_socket::scclSocket sock; // 用于网络通信的套接字
    struct scclIbCommStage stage;       // 通信阶段的状态
};

//////////////////////////////////
class scclNetIb : public scclNetBase {
public:
    // 构造函数和析构函数
    scclNetIb();
    virtual ~scclNetIb();

    // 初始化网络。
    scclResult_t init() override;
    // 返回适配器的数量。
    scclResult_t devices(int* ndev) override;
    // 获取各种设备属性。
    scclResult_t getProperties(int dev, scclNetProperties_t* props) override;
    // 创建一个接收对象并提供一个句柄以连接到它。该句柄最多可以是 SCCL_NET_HANDLE_MAXSIZE 字节，并将在排名之间交换以创建连接。
    scclResult_t listen(int dev, void* handle, void** listenComm) override;
    // 连接到一个句柄并返回一个发送 comm 对象给该对等体。
    // 此调用不应阻塞以建立连接，而应成功返回 sendComm == NULL，并期望再次调用直到 sendComm != NULL。
    scclResult_t connect(int dev, void* handle, void** sendComm) override;
    // 在远程对等体调用 connect 后最终确定连接建立。
    // 此调用不应阻塞以建立连接，而应成功返回 recvComm == NULL，并期望再次调用直到 recvComm != NULL。
    scclResult_t accept(void* listenComm, void** recvComm) override;
    // 注册/注销内存。Comm 可以是 sendComm 或 recvComm。
    // 类型是 SCCL_PTR_HOST 或 SCCL_PTR_CUDA。
    scclResult_t regMr(void* comm, void* data, int size, int type, void** mhandle) override;
    /* DMA-BUF 支持 */
    scclResult_t regMrDmaBuf(void* comm, void* data, size_t size, int type, uint64_t offset, int fd, void** mhandle) override;
    // 注销IB内存区域(MR)
    scclResult_t deregMr(void* comm, void* mhandle) override;
    // 异步发送到对等体。
    // 如果调用不能执行（或会阻塞），则可能返回 request == NULL
    scclResult_t isend(void* sendComm, void* data, int size, int tag, void* mhandle, void** request) override;
    // 异步从对等体接收。 如果调用不能执行（或会阻塞），则可能返回 request == NULL
    scclResult_t irecv(void* recvComm, int n, void** data, int* sizes, int* tags, void** mhandles, void** request) override;
    // 执行刷新/栅栏操作，以确保所有使用 SCCL_PTR_CUDA 接收到的数据对 GPU 可见
    scclResult_t iflush(void* recvComm, int n, void** data, int* sizes, void** mhandles, void** request) override;
    // 测试请求是否完成。如果 size 不为 NULL，则返回发送/接收的字节数。
    scclResult_t test(void* request, int* done, int* sizes) override;
    // 关闭并释放 send/recv comm 对象
    scclResult_t closeSend(void* sendComm) override;
    scclResult_t closeRecv(void* recvComm) override;
    scclResult_t closeListen(void* listenComm) override;

private:
    struct scclIbListenComm* ibComm = nullptr;
    // 定义一个静态变量 scclNIbDevs，用于存储 InfiniBand 设备的数量
    int scclNIbDevs = -1;

private:
    // IB异步事件处理线程主函数
    static void* scclIbAsyncThreadMain(void* args);
    // 获取IB设备的PCI路径并处理多端口和虚拟功能合并
    scclResult_t scclIbGetPciPath(char* devName, char** path, int* realPort);
    // 根据输入的宽度值，返回对应的IB(InfiniBand)链路宽度索引
    int scclIbWidth(int width);
    // 根据给定的速度值查找并返回对应的IB传输速率
    int scclIbSpeed(int speed);
    // 检查当前IB设备是否支持宽松排序(Relaxed Ordering)模式
    int scclIbRelaxedOrderingCapable(void);
    // 获取并处理用户指定的IB设备环境变量
    char* scclIbGetIbHca(int& shownIbHcaEnv, bool* searchNot, bool* searchExact);
    // 从系统文件中读取字符串内容
    scclResult_t scclGetStrFromSys(const char* path, const char* fileName, char* strValue);
    // 检查IB设备是否支持GPU Direct RDMA (GDR)
    scclResult_t scclIbGdrSupport(int ibDev);
    // 检查设备是否支持DMA-BUF功能
    scclResult_t scclIbDmaBufSupport(int dev);
    // 初始化InfiniBand Verbs资源
    scclResult_t scclIbInitVerbs(int dev, struct ibv_context* ctx, struct scclIbVerbs* verbs);
    // 创建并初始化一个InfiniBand队列对(QP)
    scclResult_t scclIbCreateQp(uint8_t ib_port, struct scclIbVerbs* verbs, int access_flags, struct ibv_qp** qp);
    // 将IB QP状态修改为RTR（Ready to Receive）状态
    scclResult_t scclIbRtrQp(struct ibv_qp* qp, uint32_t qpn, struct scclIbQpInfo* info);
    // 将IB(InfiniBand)队列对(QP)状态修改为RTS(Ready To Send)状态
    scclResult_t scclIbRtsQp(struct ibv_qp* qp);
    // 销毁IB Verbs资源
    scclResult_t scclIbDestroyVerbs(struct scclIbVerbs* verbs);
    // 从verbs请求池中获取一个未使用的请求结构体
    scclResult_t scclIbGetRequest(struct scclIbVerbs* verbs, struct scclIbRequest** req);
    // 释放IB网络请求资源。
    scclResult_t scclIbFreeRequest(struct scclIbRequest* r);
    // 执行IB网络的多发送操作
    scclResult_t scclIbMultiSend(struct scclIbSendComm* comm, int slot);
    // 通过IB Verbs RDMA写入操作向远程FIFO队列提交数据
    scclResult_t scclIbPostFifo(struct scclIbRecvComm* comm, int n, void** data, int* sizes, int* tags, void** mhandles, struct scclIbRequest* req);
};

} // namespace net_ib
} // namespace net
} // namespace hardware
} // namespace sccl