-- Copyright (c) 2024 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 lu = require 'luaunit'
local c_device_loader = require 'biz_topo.device_loader'
local c_load_info = require 'biz_topo.class.load_info'
local c_device_service = require 'device.device_service'
local cmn = require 'common'

TestDeviceLoader = {}

function TestDeviceLoader:test_parse_pcie_card_bdf_data()
    local data = '\x00\x00\x01\x01\x01\x01\x00\x00\x00\x00\x00\x01\x02\x02\x02'

    local pcie_card_info = c_device_loader.parse_pcie_card_bdf_data(data)
    -- PCIeSlot 1
    lu.assertEquals(pcie_card_info[1].bus_info.segment, 0)
    lu.assertEquals(pcie_card_info[1].bus_info.socket_id, 0)
    lu.assertEquals(pcie_card_info[1].bus_info.bus, 1)
    lu.assertEquals(pcie_card_info[1].bus_info.device, 1)
    lu.assertEquals(pcie_card_info[1].bus_info.func, 1)
    -- PCIeSlot 3
    lu.assertEquals(pcie_card_info[3].bus_info.segment, 0)
    lu.assertEquals(pcie_card_info[3].bus_info.socket_id, 1)
    lu.assertEquals(pcie_card_info[3].bus_info.bus, 2)
    lu.assertEquals(pcie_card_info[3].bus_info.device, 2)
    lu.assertEquals(pcie_card_info[3].bus_info.func, 2)
end

function TestDeviceLoader:test_get_pcie_vid_did_info()
    local bus = {
        call = function()
            return 0, '\x00\x07\xdb\x04\x01\x02\x03\x04'
        end
    }
    local pcie_info = {
        bus_info = {
            socket_id = 1,
            bus = 1,
            device = 1,
            func = 1
        }
    }
    c_device_loader.get_pcie_vid_did_info(1, bus, pcie_info)
end

function TestDeviceLoader:test_load_unload_device()
    local mock_device_load = {
        biz_topo = {
            get_mgmt_connector = function(...)
                return true, nil, {Presence = 1, LoadStatus = 0}
            end
        }
    }
    local load_info = {
        slot_id = 256
    }
    c_device_loader.load_unload_device(mock_device_load, load_info, _, _, _)
    
end

function TestDeviceLoader:test_get_id_auxid_str()
    local load_info = {
        ['id'] = 1,
        ['aux_id'] = 2,
        ['type'] = 'PCIeCard'
    }
    local device_loader = {
        handle_pcie_custom_display = function(ctx, t_id, t_aux_id)
            return t_id, t_aux_id
        end
    }
    local str_id, str_auxid = c_device_loader.get_id_auxid_str(device_loader, nil, load_info)
    lu.assertEquals(str_id, '')
    lu.assertEquals(str_auxid, '')

    load_info = {
        ['id'] = 1,
        ['aux_id'] = 2,
        ['type'] = 'OCPCard'
    }

    str_id, str_auxid = c_device_loader.get_id_auxid_str(device_loader, 'bom', load_info)
    print(str_id)
    lu.assertEquals(str_id, "00000001")
    lu.assertEquals(str_auxid, "00000002")

    load_info = {
        ['id'] = 0xfffffffe,
        ['aux_id'] = 2,
        ['type'] = 'PCIeCard'
    }
    str_id, str_auxid = c_device_loader.get_id_auxid_str(device_loader, 'bom', load_info)
    lu.assertEquals(str_id, "fffffffe")
    lu.assertEquals(str_auxid, "ffffffff")
end

function TestDeviceLoader:test_unload_device_connector()
    local c_get_instance = c_device_service.get_instance
    c_device_service.get_instance = function()
        return {
            clear_flash_checker_event = function()
            end
        }
    end
    local load_info = {
        slot_id = 256
    }
    local connector_info = {
        Id = "fffffffe",
        AuxId = "ffffffff",
    }
    local ok, _ = pcall(function()
        c_device_loader.unload_device_connector(load_info, connector_info, 1)
    end)
    lu.assertEquals(ok, true)

    c_device_service.get_instance = c_get_instance
end

function TestDeviceLoader:test_handle_pcie_custom_display()
    local records = {
        {
            Enabled = true,
            CustomType = "NPU",
            QuadrupleId = "test"
        }
    }
    local device_loader = {
        db = {
            PCIeCardDisplayCustom = {
                QuadrupleId = {
                    eq = function(id)
                        return id
                    end
                }
            },
            select = function (table)
                return {
                    where = function ()
                        return {
                            all = function ()
                                return records
                            end
                        }
                    end
                }
            end
        }
    }
    local id, aux_id = c_device_loader.handle_pcie_custom_display(device_loader, 1, 2)
    lu.assertEquals(id, "test")
    lu.assertEquals(aux_id, "NPU")
    records[1].Enabled = false
    local id, aux_id = c_device_loader.handle_pcie_custom_display(device_loader, 1, 2)
    lu.assertEquals(id, 1)
    lu.assertEquals(aux_id, 2)
end

function TestDeviceLoader:test_set_device_present_by_map()
    package.loaded['biz_topo.device_loader'] = nil
    local c_device_loader = require 'biz_topo.device_loader'

    local query_id_from_pmu = function(_sys_id, _bus, info)
        info.vid = 1
        info.did = 1
        info.sub_vid = 1
        info.sub_did = 1
    end

    local i = 0
    while true do
        i = i + 1
        local name, val = debug.getupvalue(c_device_loader.set_device_present_by_map, i)
        if not name then break end
        if name == "query_id_from_pmu" then          -- 找到目标
            debug.setupvalue(c_device_loader.set_device_present_by_map, i, query_id_from_pmu)
            break
        end
    end

    local slotid2ssbdf = { 
        [1] = {0, 0, 0x1c, 0, 0},
        [2] = {0, 0, 0x40, 0, 0}
    }

    local test_addr = {
        position = "01010101",
        update_multihost_presence = function(self, addr_multi_presence)
            self.addr_multi_presence = addr_multi_presence
        end,
        update_vid_did = function(self, vid, did)
            self.vid = vid
            self.did = did
        end,
        update_device_bdf = function(...) end,
        get_prop = function(...) return 1 end
    }

    c_device_loader.biz_topo = {
        get_pcie_addr_info = function(self, a, b)
            return true, {
                addr = test_addr
            }
        end
    }

    local load_infos
    function c_device_loader:device_load_info_persist(_sys_id, _device_type, new_pcie_load_info)
        load_infos = new_pcie_load_info
        return
    end

    function c_device_loader:task_load_unload_device(...) end

    c_device_loader:set_device_present_by_map(1, "PCIeCard", slotid2ssbdf)
    lu.assertEquals(test_addr.did,1)
    lu.assertEquals(#load_infos,2)
end

function TestDeviceLoader:test_process_pci_info()
    local connector = {
        ref_pcie_addr_info = {
            get_prop = function(_, prop)
                if prop == 'ComponentType' then
                    return 83
                end
                if prop == 'SlotID' then
                    return 1
                end
                if prop == 'SocketID' then
                    return 1
                end
                if prop == 'DevBus' then
                    return 1
                end
                if prop == 'DevDevice' then
                    return 1
                end
                if prop == 'DevFunction' then
                    return 1
                end
                if prop == 'Segment' then
                    return 1
                end
                if prop == 'MultihostPresence' then
                    return 1
                end
            end,
            update_device_bdf = function()
            end,
            update_multihost_presence = function()
            end,
            update_vid_did = function()
            end
        }
    }
    local persistent_load_info = {
        ['OCPCard'] = {
            [1] = nil
        }
    }
    local device_loader_obj = {
        hot_plug_filter = {},
        task_load_unload_device = function()
        end,
        device_load_info_persist_ocp = function()
        end,
        persistent_load_info = persistent_load_info,
        ocp_get_sub_bdf = function()
            return
        end,
        persist_ocp_systems_pcie_topo_info = function()
            return
        end
    }

    print(device_loader_obj)

    local ok, _ = pcall(function()
        c_device_loader.process_pci_info(device_loader_obj, connector, false)
        device_loader_obj.persistent_load_info = {
            ['OCPCard'] = {
                [1] = true
            }
        }
        c_device_loader.process_pci_info(device_loader_obj, connector, false)
        device_loader_obj.hot_plug_filter = {
            ['1_83'] = true
        }
        c_device_loader.process_pci_info(device_loader_obj, connector, false)
        device_loader_obj.hot_plug_filter = {
            ['1_83'] = false
        }
        c_device_loader.process_pci_info(device_loader_obj, connector, true)
        connector.ref_pcie_addr_info.get_prop = function(_, prop)
            if prop == 'ComponentType' then
                return 83
            end
            if prop == 'SlotID' then
                return 1
            end
            if prop == 'SocketID' then
                return 1
            end
            if prop == 'DevBus' then
                return 0xff
            end
            if prop == 'DevDevice' then
                return 1
            end
            if prop == 'DevFunction' then
                return 1
            end
            if prop == 'Segment' then
                return 1
            end
            if prop == 'MultihostPresence' then
                return 1
            end
        end
        c_device_loader.process_pci_info(device_loader_obj, connector, true)
    end)
    
    lu.assertEquals(ok, true)
end

function TestDeviceLoader:test_ocp_load_from_tc()
    local connector = {}
    local device_loader_obj = {
        biz_topo = {
            get_objs = function(_, prop)
                if prop == 'PCIeBizConnector' then
                    return {
                        [1] = connector
                    }
                end
            end
        },
        process_pci_info = function()
        end
    }
    local ok, _ = pcall(function()
        c_device_loader.ocp_load_from_tc(device_loader_obj)

        connector = {
            ref_pcie_addr_info = {
                get_prop = function(_, prop)
                    if prop == 'ComponentType' then
                        return 83
                    end
                    if prop == 'SlotID' then
                        return 1
                    end
                    if prop == 'SocketID' then
                        return 1
                    end
                    if prop == 'DevBus' then
                        return 1
                    end
                    if prop == 'DevDevice' then
                        return 1
                    end
                    if prop == 'DevFunction' then
                        return 1
                    end
                    if prop == 'Segment' then
                        return 1
                    end
                    if prop == 'MultihostPresence' then
                        return 1
                    end
                end,
                update_device_bdf = function()
                end,
                update_multihost_presence = function()
                end,
                update_vid_did = function()
                end
            }
        }
        c_device_loader.ocp_load_from_tc(device_loader_obj)

        connector.ref_mgmt_connector_tianchi = {
            Presence = 0
        }
        c_device_loader.ocp_load_from_tc(device_loader_obj)

        connector.ref_mgmt_connector = {
            Presence = 1
        }
        c_device_loader.ocp_load_from_tc(device_loader_obj)

        connector.ref_mgmt_connector_tianchi = {
            Presence = 1
        }
        connector.ref_mgmt_connector = {
            Presence = 0
        }
        c_device_loader.ocp_load_from_tc(device_loader_obj)
    end)
    lu.assertEquals(ok, true)
end

function TestDeviceLoader:test_device_load_info_persist_ocp()
    local persistent_load_info = {
        ['OCPCard'] = {
            [1] = {
                SlotID = 1,
                ID = 1,
                AuxID = 1,
                MultihostPresence = 1,
                save = function()
                end
            }
        }
    }
    local device_loader_obj = {
        persistent_load_info = persistent_load_info
    }
    local info = {
        id = 1,
        aux_id = 2
    }
    local ok, _ = pcall(function()
        c_device_loader.device_load_info_persist_ocp(device_loader_obj, 1, 'OCPCard', 1, info)
    end)
    lu.assertEquals(ok, true)
    lu.assertEquals(persistent_load_info['OCPCard'][1].AuxID, 2)
end

function TestDeviceLoader:test_judge_retimer_device()
    local topo_info = {
        Reserved1 = 12
    }

    c_device_loader.reset_local_db = {
        select = function()
            return {
                where = function()
                    return {
                        first = function()
                            return topo_info
                        end
                    }
                end
            }
        end,
        SystemsPCIeTopoInfo = ""
    }

    local pcie_info = {
        bus_info = {
            bus = 0
        }
    }

    local sys_id = 0
    
    lu.assertEquals(c_device_loader:is_retimer_device(pcie_info, sys_id), true)
    topo_info.Reserved1 = 0
    lu.assertEquals(c_device_loader:is_retimer_device(pcie_info, sys_id), false)
    topo_info = nil
    lu.assertEquals(c_device_loader:is_retimer_device(pcie_info, sys_id), false)
end

function TestDeviceLoader:test_get_topo_device_bdf()
    local mock_imu_cmd = {
        get_info_from_pmu = function(_, _, payload)
            if payload.address == "CLASS_CODE_INFO_ADDR" then
                if payload.bus_num == 1 then
                    return {"", "", "", "BRIDGE"}
                end
                return {"", "", "", "CLASS_CODE_NET_CARD"}
            end
            if payload.address == "BASE_ADDR_2_ADDR" then
                return {0, 2, 2, 0}
            end
        end,
        pci_info_address_list = {
            CLASS_CODE_INFO_ADDR = "CLASS_CODE_INFO_ADDR",
            BASE_ADDR_2_ADDR = "BASE_ADDR_2_ADDR"
        },
        pci_info_base_class_code = {
            CLASS_CODE_NET_CARD = "CLASS_CODE_NET_CARD"
        },
        bdf_range = {
            MIN_DEVICE = 0, MAX_DEVICE = 1,
            MIN_FUNCTION = 0, MAX_FUNCTION = 1,
        },
        addr_detail = {
            class_code_info_addr = {
                CLASS_CODE = 4
            },
            base_addr_2_addr = {
                SEC_BUS = 2,
                SUB_BUS = 3
            }
        }
    }

    package.loaded["imu"] = nil
    package.loaded["biz_topo.device_loader"] = nil

    package.loaded["imu"] = mock_imu_cmd
    print(mock_imu_cmd)

    local mock_device_loader = require ('biz_topo.device_loader')
    local sbdf = {
        socket = 1,
        bus = 1,
        device = 0,
        func = 0
    }
    local bdf = mock_device_loader.get_topo_device_bdf(1, "", sbdf)
    lu.assertEquals(bdf.bus, 2)
end

function TestDeviceLoader:test_persist_ocp_systems_pcie_topo_info()
    local info = {
        bus_info = {
            segment = 0,
            bus = 1,
            device = 0,
            func = 0
        },
        vendor_id = 114,
        device_id = 514
    }

    local property = {
        GlobalSystemId = 1,
        LocalSystemId = 0,
        Segment = 0,
        Socket = 1,
        Bus = 1,
        Device = 0,
        Funcion = 0,
        GroupId = 1,
        ComponentType = 255,
        Slot = 255,
        VendorId = 114,
        DeviceId = 514,
        SecBus = 0,
        SubBus = 0,
        Reserved1 = 0,
        CfgLowBase = 0,
        CfgLowLimit = 0,
        CfgHighBase = 0,
        CfgHighLimit = 0,
        Reserved2 = 0,
        Reserved4 = 0,
        UpdateFlag = 1,
        MMIOCnt = 0,
        MMIO0Base = 0,
        MMIO0Limit = 0,
        MMIO1Base = 0,
        MMIO1Limit = 0,
        MMIO2Base = 0,
        MMIO2Limit = 0,
        MMIO3Base = 0,
        MMIO3Limit = 0,
        MMIO4Base = 0,
        MMIO4Limit = 0,
        MMIO5Base = 0,
        MMIO5Limit = 0,
        save = function()
        end
    }

    local mock_bs = {
        new = function()
            return {
                pack = function()
                end
            }
        end
    }

    package.loaded['mc.bitstring'] = nil
    package.loaded["biz_topo.device_loader"] = nil
    package.loaded['mc.bitstring'] = mock_bs

    local mock_device_loader = require ('biz_topo.device_loader')

    mock_device_loader.reset_local_db = {
        SystemsPCIeTopoInfo = function()
            return property
        end,
        select = function()
            return {
                where = function()
                    return {
                        first = function()
                            return 
                        end
                    }
                end
            }
        end
    }

    mock_device_loader.bus = {
        pcall = function()
            return false, true
        end
    }

    local ret, _ = pcall(function()
        mock_device_loader:persist_ocp_systems_pcie_topo_info(info)
    end)
    lu.assertEquals(ret, true)
end