/*
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 manager
*/

#include "cpu_inference_manager.h"

namespace cpu_inference {

namespace {

static std::vector<int> parse_list_str(const std::string& spec)
{
    // Supports:
    // - "0,1,2,3"
    // - "0-3,8-11"
    // - "0-5:2"        -> 0,2,4
    // - "1-10:3"       -> 1,4,7,10
    // - "0-3,5-9:2,12" -> 0,1,2,3,5,7,9,12
    if (spec.empty()) return {};
    static const std::regex valid_format(
        R"(^\d+(?:-\d+(?::\d+)?)?(?:,\d+(?:-\d+(?::\d+)?)?)*$)"
    );
    if (!std::regex_match(spec, valid_format)) {
        throw std::invalid_argument("Invalid format for CPU affinity spec: " + spec);
    }

    std::set<int> out;
    std::stringstream ss(spec);
    std::string part;
    while (std::getline(ss, part, ',')) {
        // trim spaces
        part.erase(0, part.find_first_not_of(" \t\n\r"));
        part.erase(part.find_last_not_of(" \t\n\r") + 1);
        if (part.empty()) continue;

        int step = 1;
        std::string range_part = part;
        const auto colon_pos = part.find(':');
        if (colon_pos != std::string::npos) {
            range_part = part.substr(0, colon_pos);
            const std::string step_str = part.substr(colon_pos + 1);
            step = std::stoi(step_str);
            if (step <= 0) {
                throw std::invalid_argument("Invalid step in CUSTOM_CPU_AFFINITY: " + part);
            }
        }

        const auto dash_pos = range_part.find('-');
        if (dash_pos != std::string::npos) {
            const int start = std::stoi(range_part.substr(0, dash_pos));
            const int end = std::stoi(range_part.substr(dash_pos + 1));
            if (start > end) {
                throw std::invalid_argument("Invalid range in CUSTOM_CPU_AFFINITY: " + range_part);
            }
            for (int v = start; v <= end; v += step) {
                out.insert(v);
            }
        } else {
            out.insert(std::stoi(range_part));
        }
    }

    return std::vector<int>(out.begin(), out.end());
}

static std::map<int, std::vector<int>> parse_numa_info(const std::vector<int>& cpu_ids)
{
    std::map<int, std::vector<int>> cpu_affinity;
    for (int cpu_id : cpu_ids) {
        const int node = numa_node_of_cpu(cpu_id);
        if (node >= 0) {
            cpu_affinity[node].push_back(cpu_id);
        }
    }
    for(auto it = cpu_affinity.begin(); it != cpu_affinity.end(); ++it) {
        std::sort(it->second.begin(), it->second.end());
    }
    return cpu_affinity;
};

// 初始化进程亲和性
void init_process_affinity(const std::vector<int>& cpu_ids) {
    if (cpu_ids.size() == 0) {
        return;
    }
    int num_threads = cpu_ids.size();
    // 设置线程数量
    omp_set_num_threads(num_threads);
    omp_set_dynamic(0); 
    
    // 绑核
    #pragma omp parallel
    {
        int current_tid = omp_get_thread_num();

        cpu_set_t mask;
        CPU_ZERO(&mask);
        CPU_SET(cpu_ids[current_tid], &mask);

        int ret = sched_setaffinity(0, sizeof(cpu_set_t), &mask);
        if (ret == -1) {
            std::cerr << "Error! Binding core failed!\n";
            exit(0);
        }
    }
}

} // namespace

// 初始化
void CPUInferenceManager::init(py::dict run_config, py::dict model_config){
    try{
        std::string model_type = model_config["model_type"].cast<std::string>();
        std::string quant_type = run_config["quant_type"].cast<std::string>();
        std::string cpu_affinity = run_config["cpu_affinity"].cast<std::string>();
        int nrc = 2;
        if(run_config.contains("nrc")){
            nrc = std::stoi(run_config["nrc"].cast<std::string>());
        }
        set_cpu_affinity(cpu_affinity);
        set_quant_type(quant_type);
        if(model_type == "qwen3_moe"){
            Qwen3MoeConfig config;
            config.num_hidden_layers = model_config["num_hidden_layers"].cast<int>();
            config.hidden_size = model_config["hidden_size"].cast<int>();
            config.moe_intermediate_size = model_config["moe_intermediate_size"].cast<int>();
            config.num_experts = model_config["num_experts"].cast<int>();
            config.num_experts_per_tok = model_config["num_experts_per_tok"].cast<int>();
            this->model = new Qwen3MoeModel(this->numa_list, this->weight_quant, 
                this->input_quant, nrc, config);
        }
        else{
            throw std::runtime_error(std::string("not supported model type") + model_type);
        }

        this->nrc = nrc;
        this->model_type = model_type;
        init_process_affinity(this->cpu_ids);      
        if(this->model_weight){
            delete this->model_weight;
            this->model_weight = nullptr;
        }
        this->model_weight = new ModelWeightBase(this->model_type, this->weight_quant, this->cpu_affinity);
    }
    catch (const pybind11::error_already_set& e) {
        // 捕获 Python 异常（如 KeyError, TypeError 来自 .cast 失败等）
        // 自动包含 traceback，会原样抛出到 Python 层
        throw;
    }
    catch (const std::exception& e) {
        // 捕获所有标准 C++ 异常（包括 runtime_error, invalid_argument 等）
        throw std::runtime_error(std::string("Initialization failed: ") + e.what());
    }
    catch (...) {
        // 捕获未知异常
        throw std::runtime_error("Unknown error occurred during CPUInferenceManager initialization.");
    }
}

void CPUInferenceManager::set_quant_type(std::string quant_type)
{
    if(quant_type == "fp16fp16"){
        this->weight_quant = "fp16";
        this->input_quant = "fp16";
    }
    else if(quant_type == "q4q8"){
        this->weight_quant = "q4_0";
        this->input_quant = "q8_0";
    }
    else if(quant_type == "q8q8"){
        this->weight_quant = "q8_0";
        this->input_quant = "q8_0";
    }
    else if(quant_type == "q8align"){
        this->weight_quant = "q8align";
        this->input_quant = "q8align";
    }
    else {
        this->weight_quant = "fp16";
        this->input_quant = "fp16";
        std::cerr << "Invalid quant type, use default quant type : fp16fp16" << std::endl; 
    }
}

void CPUInferenceManager::set_model_type(std::string model_type)
{
    this->model_type = model_type;
}

void CPUInferenceManager::set_cpu_affinity(std::string affinity_env)
{
    try {
        this->cpu_ids = parse_list_str(std::string(affinity_env));
    } catch (const std::exception& e) {
        std::cerr << "Error! Failed to parse CUSTOM_CPU_AFFINITY='" << affinity_env
                << "': " << e.what() << std::endl;
        return;
    }
    CHECK_MSG(!(this->cpu_ids.empty()), "cpu_ids is empty");
    this->cpu_affinity = parse_numa_info(this->cpu_ids);
    this->numa_list.clear();
    for(const auto& numa_cpu_item : this->cpu_affinity){
        this->numa_list.push_back(numa_cpu_item.first);
    }
}

void CPUInferenceManager::load_model(py::list weight_infos, bool is_load_over){
    for (const auto& weight_info : weight_infos) {
        py::dict weight_info_dict = weight_info.cast<py::dict>();
        py::object weight = weight_info_dict["data"];
        py::object data_ptr_obj = weight.attr("data_ptr")();
        void* data_ptr = reinterpret_cast<void*>(data_ptr_obj.cast<uint64_t>());
        py::tuple py_shape = weight.attr("shape").cast<py::tuple>();
        std::vector<int> shape;
        std::string dtype = "float16";
        for (size_t i = 0; i < py_shape.size(); ++i) {
            shape.push_back(py_shape[i].cast<int64_t>());
        }
        std::string name = weight_info_dict["meta"].cast<std::string>();
        this->model_weight->add_weight(name, dtype, shape, data_ptr);
    }

    if(is_load_over){
        this->model->load_model(this->model_weight);
    }
}

void CPUInferenceManager::forward(torch::Tensor& expert_output, 
    const torch::Tensor& hidden_states, 
    const torch::Tensor& router_logits, 
    int64_t layer_id){
    this->model->forward(expert_output, hidden_states, router_logits, layer_id);
}

}; // namespace cpu_inference