#pragma once

#include <stdint.h>
#include "base.h"

namespace sccl {
namespace hardware {
namespace net {

typedef enum {
    SCCL_PTR_HOST   = 0x1,
    SCCL_PTR_CUDA   = 0x2,
    SCCL_PTR_DMABUF = 0x4
} sccl_ptr_t;

constexpr int SCCL_SOCKET_SEND        = 0;
constexpr int SCCL_SOCKET_RECV        = 1;
constexpr int SCCL_NET_HANDLE_MAXSIZE = 128;

////////////////////////////////// 用于定义网络设备 //////////////////////////////////
typedef struct {
    char* name;     // 主要用于日志记录。
    char* pciPath;  // PCI设备在/sys中的路径。
    uint64_t guid;  // NIC芯片的唯一标识符。对于具有多个PCI功能（物理或虚拟）的卡非常重要。
    int ptrSupport; // [SCCL_PTR_HOST|SCCL_PTR_CUDA|SCCL_PTR_DMABUF]
    int speed;      // 端口速度，单位为Mbps。
    int port;       // 端口号。
    float latency;  // 网络延迟
    int maxComms;   // 我们可以创建的最大通信数量
    int maxRecvs;   // 最大分组接收数量。
} scclNetProperties_t;

/**
 * @brief scclNetBase 类定义了网络通信的基础接口
 *
 * 该类是一个抽象基类，提供了网络初始化、设备管理、连接建立、内存注册、
 * 数据传输等核心功能的纯虚函数接口。具体实现应由派生类完成。
 *
 * 主要功能包括：
 * - 网络初始化和设备属性查询
 * - 监听/连接建立和管理
 * - 内存注册和注销
 * - 异步发送/接收操作
 * - 请求状态测试
 * - 连接关闭
 *
 * 接口设计为非阻塞式，支持异步操作。
 */
typedef class scclNetBase {
public:
    // 构造函数和析构函数
    scclNetBase(const char* net_name) : name(net_name) {};
    virtual ~scclNetBase() {};

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

public:
    // 网络的名称（主要用于日志）
    const char* name;

} scclNet_t;

////////////////////////////////// 功能函数 //////////////////////////////////
// 初始化 ROCm 库
scclResult_t rocmLibraryInit(void);
struct netIf {       // 网络接口结构体
    char prefix[64]; // 网络前缀
    int port;        // 端口号
};

// 解析字符串列表，将结果存储在网络接口列表中
int parseStringList(const char* string, struct netIf* ifList, int maxList);

// 根据给定的字符串和端口，匹配网络接口列表中的接口
bool matchIfList(const char* string, int port, struct netIf* ifList, int listSize, bool matchExact);

// 打印网络属性信息
scclResult_t printNetProps(const scclNetProperties_t* props, int rank, int localRank);

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