/*
Copyright (c) 2025-2025 Huawei Technologies Co., Ltd.

sysHAX-adapter is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2.
You may obtain a copy of Mulan PSL v2 at:
    http://license.coscl.org.cn/MulanPSL2
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
PURPOSE.
See the Mulan PSL v2 for more details.
Created: 2026-1-31
Desc: CPU inference memory manager
*/

#include "memory_manager.h"

namespace cpu_inference {

MemoryManager::MemoryManager() {
    if (numa_available() == -1) {
        std::cerr << "[MemoryManager] NUMA not available, assuming 1 node.\n";
        tot_numa_cnt = 1;
    } else {
        numa_set_strict(1);
        tot_numa_cnt = numa_max_node() + 1;
    }
}

MemoryManager::~MemoryManager() {}

MemoryManager& MemoryManager::get() {
    static MemoryManager instance;
    return instance;
}

void MemoryManager::clear_all(){
    alloc_memory.clear();
}

Tensor& MemoryManager::alloc_single_numa(const std::string& name,
                                            const std::vector<int>& shape,
                                            int numa_node,
                                            const std::string& dtype,
                                            const std::string& quant_type
                                            ) {
    if (numa_node < 0 || numa_node >= tot_numa_cnt) {
        throw std::out_of_range("Invalid NUMA node: " + std::to_string(numa_node));
    }

    // 单 NUMA：直接构造
    auto res = alloc_memory.emplace(
        name,
        Tensor(name, {numa_node}, shape, dtype, quant_type)
    );
    if (!res.second) {
        throw std::runtime_error("Tensor name already exists: " + name);
    }
    return res.first->second;
}

Tensor& MemoryManager::alloc_multi_numa(const std::string& name,
                                          const std::vector<int>& shape,
                                          const std::vector<int>& numas,
                                          const std::string& dtype,
                                          const std::string& quant_type
                                          ) {
    if (numas.empty()) {
        throw std::invalid_argument("numas list is empty");
    }

    // 去重并验证 NUMA 节点有效性
    std::vector<int> unique_numas(numas.begin(), numas.end());
    std::sort(unique_numas.begin(), unique_numas.end());
    unique_numas.erase(std::unique(unique_numas.begin(), unique_numas.end()), unique_numas.end());

    for (int node : unique_numas) {
        if (node < 0 || node >= tot_numa_cnt) {
            throw std::out_of_range("Invalid NUMA node: " + std::to_string(node));
        }
    }

    auto res = alloc_memory.emplace(
        name,
        Tensor(name, numas, shape, dtype, quant_type)
    );

    return res.first->second;
}

Tensor& MemoryManager::alloc_all_numa(const std::string& name,
                                        const std::vector<int>& shape,
                                        const std::string& dtype,
                                        const std::string& quant_type) {
    std::vector<int> all_numas(tot_numa_cnt);
    std::iota(all_numas.begin(), all_numas.end(), 0);
    return alloc_multi_numa(name, shape, all_numas, dtype, quant_type);
}

Tensor& MemoryManager::get_memory(const std::string& name) {
    auto it = alloc_memory.find(name);
    if (it == alloc_memory.end()) {
        throw std::out_of_range("Tensor not found: " + name);
    }
    return it->second;
}

int MemoryManager::get_tot_numa_cnt() const {
    return tot_numa_cnt;
}

} // namespace cpu_inference