#pragma once

#include <string.h>
#include "base.h"
#include "bootstrap.h"

namespace sccl {
namespace hardware {
namespace topology {
namespace graph {
typedef bootstrap::physical_links::scclTopoNode_t scclTopoNode_t;
typedef bootstrap::scclNodeInfo_t scclNodeInfo_t;
typedef bootstrap::BootstrapComm_t BootstrapComm_t;
typedef topology::bootstrap::Bootstrap Bootstrap;

// 定义 topoPathType_t 枚举类型，用于表示不同的路径类型。
typedef enum topoPathType {
    PATH_LOC = 0, // 本地路径
    PATH_NVL = 1, // 通过 NVLink 连接
    PATH_NVB = 2, // 通过中间 GPU 使用 NVLink 连接
    PATH_PIX = 3, // 通过最多一个 PCIe 桥连接
    PATH_PXB = 4, // 通过多个 PCIe 桥连接（不经过 PCIe 主桥）
    PATH_PXN = 5, // GPU 和 NIC 之间通过中间 GPU 连接
    PATH_PHB = 6, // 通过 PCIe 以及 PCIe 主桥连接
    PATH_SYS = 7, // 通过 PCIe 以及 NUMA 节点之间的 SMP 互连连接
    PATH_NET = 8, // 通过网络连接
    PATH_DIS = 9  // 断开连接
} topoPathType_t;

// GPU 连接其他GPU硬件的直连类型
typedef enum LinkType : uint8_t {
    LINK_NONE = 0, // 本地路径
    LINK_LOC  = 1, // 本地路径
    LINK_NVL  = 2, // 通过 NVLink 连接
    LINK_PIX  = 3, // 通过 PCIe 桥连接
    LINK_PXN  = 4, // GPU 和 GPU 之间通过中间 NIC 连接，包括 PCIe 主桥
    LINK_NET  = 5  // 通过网络连接
} LinkType_t;

typedef struct scclTopoGraph {
    scclTopoGraph() = delete; // 删除默认构造函数
    scclTopoGraph(int nRanks) : nRanks(nRanks), transport_map(nullptr, 0) {
        // 分配transport_map的内存
        uint8_t* raw_transport_map = static_cast<uint8_t*>(calloc(nRanks * nRanks, sizeof(uint8_t)));
        if(raw_transport_map == nullptr) {
            // 处理内存分配失败的情况
            throw std::bad_alloc();
        }
        // 使用ByteSpanArray初始化transport_map
        transport_map = ByteSpanArray<uint8_t>(raw_transport_map, nRanks * nRanks);
    }
    virtual ~scclTopoGraph() {
        // 释放transport_map的内存
        free(transport_map.data());
    }

    uint8_t* getTransportMapRowStart(int row) { return transport_map[row * nRanks]; }
    uint8_t* getTransportMapData(int row, int col) { return transport_map[row * nRanks + col]; }

    // 打印transport_map
    scclResult_t printTransportMap() {
        for(int i = 0; i < this->nRanks; ++i) {
            for(int j = 0; j < this->nRanks; ++j) {
                uint8_t* value = this->getTransportMapData(i, j);
                if(value != nullptr) {
                    printf("%d ", *value);
                } else {
                    printf("nullptr ");
                }
            }
            printf("\n");
        }
        return scclSuccess;
    }

public:
    // 使用无序映射存储图的有效节点
    std::unordered_map<uint64_t, scclTopoNode_t> graph_nodes;
    // 使用无序映射存储从每个GPU节点到其他GPU节点的所有路径，[start_node_id][end_node_id] = {path1, path2}
    std::unordered_map<uint64_t, std::unordered_map<uint64_t, std::vector<std::vector<uint64_t>>>> gpu_paths;

    // 传输位图
    ByteSpanArray<uint8_t> transport_map; // 使用ByteSpanArray存储transport_map
    int nRanks;                           // 记录GPU节点的数量
} scclTopoGraph_t;

} // namespace graph
} // namespace topology
} // namespace hardware
} // namespace sccl