-- Copyright (c) 2025 Huawei Technologies Co., Ltd.
-- openUBMC 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.

local CPU_INTERFACE<const> = 'bmc.kepler.Systems.Processor.CPU'
local PROCESSOR_INTERFACE<const> = 'bmc.kepler.Systems.Processor'

local m = {}

function m.get_logical_physical_tab(cpu_list)
    if not cpu_list or next(cpu_list) == nil then
        return
    end
    local ok, obj
    local system_id, logical_id, physical_id
    local tab = {}
    for _, path in ipairs(cpu_list) do
        -- 获取系统ID
        ok = pcall(function()
            obj = mdb.get_object(bus, path, PROCESSOR_INTERFACE)
        end)
        if ok then
            system_id = obj.SystemId
        else
            goto continue
        end
        -- 获取物理、逻辑ID
        ok = pcall(function()
            obj = mdb.get_object(bus, path, CPU_INTERFACE)
        end)
        if ok then
            logical_id, physical_id = obj.LogicalId, obj.PhysicalId
        else
            goto continue
        end
        -- 按照系统维度保存物理、逻辑ID
        if not tab[system_id] then
            tab[system_id] = {}
        end
        tab[system_id][#(tab[system_id]) + 1] = {LogicalId = logical_id, PhysicalId = physical_id}
        ::continue::
    end
    for _, v in pairs(tab) do
        table.sort(v, function(a, b) return a['LogicalId'] < b['LogicalId'] end)
    end
    return tab
end

function m.get_associate_resource_by_socket_and_sys(sys_tab, pcie_socket_id, cpu_list)
    -- sys_tab为空按system 1 处理
    if not sys_tab or next(sys_tab) == nil then
        sys_tab = {1}
    end
    local cpu_tab = m.get_logical_physical_tab(cpu_list)
    if not cpu_tab or next(cpu_tab) == nil then
        return {}
    end
    for _, sys_id in pairs(sys_tab) do
        if not cpu_tab[sys_id] then
            goto continue
        end
        if pcie_socket_id == 0xff then
            return 'PCIeSwitch'
        elseif pcie_socket_id == 0xfe then
            return 'PCH'
        -- 取前两个cpu
        elseif pcie_socket_id == 0xfd then
            return 'CPU' .. cpu_tab[sys_id][1].PhysicalId .. ',CPU' .. cpu_tab[sys_id][2].PhysicalId
        -- 取前四个cpu
        elseif pcie_socket_id == 0xfc then
            return 'CPU' .. cpu_tab[sys_id][1].PhysicalId .. ',CPU' .. cpu_tab[sys_id][2].PhysicalId .. 
                    ',CPU' .. cpu_tab[sys_id][3].PhysicalId .. ',CPU' .. cpu_tab[sys_id][4].PhysicalId
        -- 匹配唯一cpu
        else
            for _, v in pairs(cpu_tab[sys_id]) do
                if v.LogicalId == pcie_socket_id then
                    return 'CPU' .. v.PhysicalId
                end
            end
        end
        ::continue::
    end
end

function m.get_pcie_socket_id(objectid, pcie_paths)
    local pos = objectid[4]

    for _, v in ipairs(pcie_paths) do
        if pos == mdb.get_object(bus, v, 'bmc.kepler.Object.Properties').ObjectIdentifier[4] then
            return mdb.get_object(bus, v, 'bmc.kepler.Systems.PCIeDevices.PCIeDevice')['SocketID']
        end
    end

    return nil
end

function m.get_associate_resource_by_socket_and_sys_network_adapter(system_id, socketid, cpu_list)
    local sys_tab = {system_id}
    return m.get_associate_resource_by_socket_and_sys(sys_tab, socketid, cpu_list)
end

function m.get_multi_associate_resource_by_socket_and_sys_network_adapter(system_id, socketids, cpu_list, associate_resource)
    local sys_tab = {system_id}
    local resources = {}
    for k, v in pairs(socketids) do
        resources[#resources + 1] = m.get_associate_resource_by_socket_and_sys(sys_tab, v, cpu_list)
    end
    if #resources == 0 then
        resources[1] = associate_resource
    end
    return resources
end

function m.get_multi_associate_resource_by_socket_and_sys_pcie_device(system_id, device_path, cpu_list, associate_resource)
    local socket_ids = {}
    local sub_objects = mdb.get_sub_objects(bus, device_path .. '/PCIeFunctions/', 'bmc.kepler.Systems.PCIeDevice.PCIeFunction')
    for _, obj in pairs(sub_objects) do
        -- 以下代码可以保证两个功能为同一个CPU时不会被加载第二次，同时避免因为键为0导致无法遍历
        socket_ids[obj.RelatedProcessorId + 1] = obj.RelatedProcessorId
    end

    local sys_tab = {system_id}
    local resources = {}
    for k, v in pairs(socket_ids) do
        resources[#resources + 1] = m.get_associate_resource_by_socket_and_sys(sys_tab, v, cpu_list)
    end

    if #resources == 0 then
        resources[1] = associate_resource
    end
    
    return resources
end

function m.get_drive_associate_resource_by_resourceid(system_id, resourceid, cpu_list)
    if not (resourceid > 0 and resourceid < 255) then
        return nil
    end

    -- 实际socket id为resourceid - 1
    local socketid = resourceid - 1
    local sys_tab = {system_id}

    if resourceid == 0xfe then
        socketid = resourceid
    end

    return m.get_associate_resource_by_socket_and_sys(sys_tab, socketid, cpu_list)
end

return m