#include <string.h>
#include "graph_utils.h"

namespace sccl {
namespace hardware {
namespace topology {
namespace graph {

scclTopoGraph::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);
}

scclTopoGraph::~scclTopoGraph() {
    // 释放transport_map的内存
    free(transport_map.data());
}

// 打印transport_map
scclResult_t scclTopoGraph::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;
}

// 打印gpu_paths信息的函数
scclResult_t scclTopoGraph::printGPUPaths() {
    for(const auto& start_pair : gpu_paths) {
        uint64_t start_node_id = start_pair.first;
        auto start_node_it     = graph_nodes.find(start_node_id);
        if(start_node_it != graph_nodes.end()) {
            std::cout << "Paths starting from node: ";
            start_node_it->second.printNodeInfo("Start Node");
        } else {
            std::cout << "Start node ID " << start_node_id << " not found in graph nodes." << std::endl;
            continue;
        }

        for(const auto& end_pair : start_pair.second) {
            uint64_t end_node_id = end_pair.first;
            auto end_node_it     = graph_nodes.find(end_node_id);
            if(end_node_it != graph_nodes.end()) {
                std::cout << "  to node: ";
                end_node_it->second.printNodeInfo("End Node");
            } else {
                std::cout << "  End node ID " << end_node_id << " not found in graph nodes." << std::endl;
                continue;
            }

            std::cout << "  Paths:" << std::endl;
            for(const auto& path : end_pair.second) {
                std::cout << "    Path: ";
                for(const auto& node_id : path) {
                    auto node_it = graph_nodes.find(node_id);
                    if(node_it != graph_nodes.end()) {
                        node_it->second.printNodeInfo("    ");
                    } else {
                        std::cout << "    Node ID " << node_id << " not found in graph nodes." << std::endl;
                    }
                }
                std::cout << std::endl;
            }
        }
    }
    return scclSuccess;
}

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