-- 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 lu = require 'luaunit'
local nvme_admin_command_process = require 'nvme.nvme_mi_protocol.nvme_admin_command_process'
local nvme_admin_command = require 'nvme.nvme_mi_protocol.nvme_mi_admin_command'
local nvme_mctp = require 'nvme.nvme_mi_protocol.nvme_mi_mctp'
local NVME_MI_OPTIONS = require 'nvme.nvme_mi_protocol.nvme_mi_def'
local common_def = require 'common_def'
local nvme_obj = require 'nvme.nvme_object'
local storage_bus = require 'storage_bus'
local mctp_lib = require 'mctp_lib'
local c_storageconfig = require 'storageconfig.storageconfig_object'
local log = require 'mc.logging'

TESTNvmeObject = {}

-- 测试nvme盘是否支持hw defined smart log
function TESTNvmeObject:test_check_support_hw_defined_smart_log()
    local ret = nvme_obj:check_support_hw_defined_smart_log()
    lu.assertEquals(ret, common_def.INVALID_U8)

    nvme_obj.nvme_mi_mctp_obj = {
        queue = function (func)
            func()
        end,
        nvme = nvme_obj,
        nvme_mi_obj = {
            ReadCtrlList = function ()
                return {
                    value = function ()
                        return 1
                    end
                }
            end
        }
    }
    nvme_admin_command_process.check_support_hw_defined_without_uuid = function ()
        return true
    end
    ret = nvme_obj:check_support_hw_defined_smart_log()
    lu.assertEquals(ret, NVME_MI_OPTIONS.SUPPORT_HW_DEFINED_WITHOUT_UUID)

    nvme_admin_command_process.check_support_uuid_by_identify_ctrl = function ()
        return true
    end
    nvme_admin_command_process.check_support_hw_defined_without_uuid = function ()
        return false
    end
    ret = nvme_obj:check_support_hw_defined_smart_log()
    lu.assertEquals(ret, 0)

    nvme_admin_command_process.check_specified_opcode_support_uuid = function ()
        return true
    end
    ret = nvme_obj:check_support_hw_defined_smart_log()
    lu.assertEquals(ret, 0)

    nvme_admin_command_process.get_specified_uuid_index = function ()
        return 2
    end
    ret = nvme_obj:check_support_hw_defined_smart_log()
    lu.assertEquals(ret, NVME_MI_OPTIONS.SUPPORT_HW_DEFINED_WITH_UUID)
    lu.assertEquals(nvme_obj.uuid_index, 2)

    nvme_obj.nvme_mi_mctp_obj = nil
end

-- 测试nvme盘获取hw defined smart log
function TESTNvmeObject:test_get_nvme_hw_defined_smart_info()
    local ret = nvme_obj:get_nvme_hw_defined_smart_info()
    lu.assertEquals(ret, nil)

    nvme_obj.nvme_mi_mctp_obj = {ctrl_id = 1, queue = function () end, nvme = { Slot = 0 }}
    ret = nvme_obj:get_nvme_hw_defined_smart_info()
    lu.assertEquals(ret, nil)

    nvme_obj.support_hw_defined_smart_log = NVME_MI_OPTIONS.SUPPORT_HW_DEFINED_WITHOUT_UUID
    nvme_admin_command.get_hw_defined_smart_log = function ()
        return '\x5c\x23\x00\x5f\x64\x5f\x00\x00\x00\x00\x00\x00' .. -- 0x5c
            '\x69\x23\x00\x62\x64\x62\x00\x00\x00\x00\x00\x00' .. -- 0x69
            '\x5b\x32\x00\x64\x64\xfa\x01\x00\x02\xff\x01\x00' .. -- 0x5b
            '\x68\x32\x00\x64\x64\xdd\x05\xe8\x05\xe3\x05\x00' .. -- 0x68
            '\x67\x32\x00\x64\x64\x90\x5f\x01\x00\x00\x00\x00' .. -- 0x67
            '\x5a\x32\x00\x64\x64\xf8\x2a\x00\x00\x00\x00\x00' .. -- 0x5a
            '\x31\x32\x00\x64\x64\x03\x01\x00\x00\x00\x00\x00' .. -- 0x31
            '\x3a\x32\x00\x64\x64\x0d\x00\x00\x00\x00\x00\x00' .. -- 0x3a
            '\x33\x33\x00\x73\x73\x03\x01\x00\x00\x00\x00\x56' .. -- 0x33
            '\x3c\x33\x00\xa2\xa2\x0d\x00\x00\x00\x00\x00\x3d' .. -- 0x3c
            '\x52\x32\x00\x64\x64\x00\xca\xc5\x94\x13\x04\x00' .. -- 0x52
            '\x53\x32\x00\x64\x64\x00\x00\x00\x00\x00\x00\x00' .. -- 0x53
            '\x4e\x23\x00\x59\x64\x23\x23\x00\x00\x00\x00\x00' .. -- 0x4e
            '\x4f\x23\x00\x59\x64\x23\x23\x00\x00\x00\x00\x00'    -- 0x4f
    end
    local expected_ret = {
        tlc_spare_block	= 100,
        slc_spare_block = 100,
        remnant_wearout = 95,
        slc_avg_ec = 1507,
        slc_used_lifespan =	'89.95',
        tlc_avg_ec = 511,
        hw_defined_nand_write_h = 0,
        tlc_pe_cycle = 11000,
        slc_pe_cycle = 90000,
        tlc_used_lifespan =	'89.95',
        hw_defined_nand_write_l	= 4482146880000,
    }
    ret = nvme_obj:get_nvme_hw_defined_smart_info()
    for k, v in pairs(expected_ret) do
        lu.assertEquals(v, ret[k])
    end
end

-- 测试nvme盘init mctp
function TESTNvmeObject:test_init_nvme_mi_mctp()
    local drive = {write_amplification_info = {}}
    nvme_obj.protocol = 1
    local origin_func = nvme_obj.check_support_mctp
    nvme_obj.check_support_mctp = function ()
        return true
    end
    nvme_obj.new_task = function(nvme_obj)
        return {
            loop = function (self, func)
                local task = {
                    stop = function ()
                    end,
                    is_exit = true
                }
                func(task)
                return {
                    set_timeout_ms = function()
                    end
                }
            end
        }
    end
    nvme_obj:init_nvme_mi_mctp(drive)
    lu.assertNotEquals(nvme_obj.nvme_mi_mctp_obj, nil)
    nvme_obj.nvme_mi_mctp_obj = nil
    nvme_obj.check_support_mctp = origin_func
end

-- 测试文件安全函数调用
function TESTNvmeObject:test_collect_nvme_log()
    local drive = {write_amplification_info = {}}
    nvme_obj.protocol = 1
    local origin_func = nvme_obj.check_support_mctp
    nvme_obj.check_support_mctp = function ()
        return true
    end
    local func = nvme_mctp.create_mctp_endpoint_for_smbus
    nvme_mctp.create_mctp_endpoint_for_smbus = function ()
        return true
    end
    nvme_obj.get_position = function ()
        return ""
    end
    nvme_obj:init_nvme_mi_mctp(drive)
    lu.assertNotEquals(nvme_obj.nvme_mi_mctp_obj, nil)
    ret = pcall(nvme_obj.nvme_mi_mctp_obj.collect_nvme_log, nvme_obj.nvme_mi_mctp_obj)
    lu.assertEquals(ret, true)
    nvme_obj.nvme_mi_mctp_obj = nil
    nvme_obj.check_support_mctp = origin_func
    nvme_mctp.create_mctp_endpoint_for_smbus = func
end

-- 测试NVMe盘smart日志收集
function TESTNvmeObject:test_dump_nvme_smart_log()
    local obj = {
        nvme_mi_mctp_obj = {
            collect_nvme_smart_log = function(obj) end
        }
    }
    ret0 = pcall(nvme_obj.dump_nvme_smart_log, obj, true)
    lu.assertEquals(ret0, true)
    ret1 = pcall(nvme_obj.dump_nvme_smart_log, obj, false)
    lu.assertEquals(ret1, true)
end

function TESTNvmeObject:test_create_smbus_endpoint()
    local drive = {}
    nvme_obj.protocol = 1
    nvme_obj.is_init_nvmemi_mctp = false
    nvme_obj.nvme_mi_mctp_obj = nil
    local origin_func = nvme_obj.check_support_mctp
    nvme_obj.check_support_mctp = function ()
        return true
    end
    storage_bus.get_instance = function ()
        return { bus = true }
    end
    local mctp_fun = mctp_lib.get_endpoint_and_transport
    mctp_lib.get_endpoint_and_transport = function ()
        return {}, {}
    end
    nvme_obj:init_nvme_mi_mctp(drive)
    lu.assertNotEquals(nvme_obj.nvme_mi_mctp_obj, nil)
    nvme_obj.nvme_mi_mctp_obj = nil
    nvme_obj.check_support_mctp = origin_func
    mctp_lib.get_endpoint_and_transport = mctp_fun
end

function test_vpd_nvme_mi_get_vender_id()
    -- 模拟 common_def 和外部依赖
    local common_def = {
        RECORD_ITEM_PRODUCT_NAME = "PM9D3a",
        INVALID_STRING = "INVALID",
        NVME_VPD_VENDOR_ID_FOR_INTEL_P4600 = "0x80860001",
        NVME_VPD_VENDOR_ID_FOR_INTEL_P5800X = "0x80860002",
        NVME_VPD_VENDOR_ID_FOR_INTEL = "0x8086",
        NVME_VPD_VENDOR_ID_FOR_SAMSUNG_PM9A3 = "0x144d0005",
        NVME_VPD_VENDOR_ID_FOR_SAMSUNG_PM9D3A = "0x144da80e",
        NVME_VPD_VENDOR_ID_FOR_MEMBLAZE_7940 = "0x1c5f0007",
        MANUFACTURE_ID_MAP = {
            ['HUAWEI'] = 0x19e5,
            ['INTEL'] = 0x8086,
            ['WESTDIGI'] = 0x1b96,
            ['RAMAXEL'] = 0x1e81,
            ['DAPUSTOR'] = 0x1e3b,
            ['UMIS'] = 0x1cc4,
            ['KIOXIA'] = 0x1e0f,
            ['SAMSUNG'] = 0x144d,
            ['STARBLZ'] = 0x9d32,
            ['MEMBLAZE'] = 0x1c5f,
            ['SHANNON'] = 0x1cb0,
            ['MICRON'] = 0x1344,
            ['DERA'] = 0x1d78,
            ['SOLIDIGM'] = 0x025e,
            ['unknown'] = 0xFFFF
        }
    }

    -- 修改返回的产品名称为 SAMSUNG
    local result = nvme_obj:vpd_nvme_mi_get_vender_id("SAMSUNG")
    lu.assertEquals(result,common_def.MANUFACTURE_ID_MAP["SAMSUNG"])

    -- 修改返回的产品名称为 SOLIDIGM
    local result = nvme_obj:vpd_nvme_mi_get_vender_id("SOLIDIGM")
    lu.assertEquals(result,common_def.MANUFACTURE_ID_MAP["SOLIDIGM"])

    -- 修改返回的产品名称为 huawei
    result = nvme_obj:vpd_nvme_mi_get_vender_id("HUAWEI")
    lu.assertEquals(result,common_def.MANUFACTURE_ID_MAP["HUAWEI"])
end

function test_get_nvme_smart_info()
    local result = nvme_obj:get_nvme_smart_info()
    lu.assertEquals(result, nil)
end

-- 测试NVMe盘更新Revision
function TESTNvmeObject:test_get_nvme_revision()
    local obj = {
        SSDChip = {
            Read = function ()
                return 'h'
            end
        },
        ManufacturerId = 0xffffffff,
        Manufacturer = 'N/A',
        vpd_nvme_mi_get_product_info = function ()
            return 'SAMSUNG'
        end,
        vpd_nvme_mi_get_vender_id = function ()
        end
    }
    local revision = nvme_obj.get_firmeware_version(obj)
    lu.assertEquals(revision, 'hhhhhhhh')
end

function TESTNvmeObject:test_nvme_init()
    nvme_obj.super = {
        init = function()
        end
    }
    local obj = {
        update_nvme_endpoint = function ()
        end,
        on_smbios_status_changed = {
            on = function (self, cb)
                cb()
            end
        },
        new_task = function ()
            task = {
                stop = function ()
                end
            }
            return {
                loop = function (self, cb)
                    cb(task)
                end
            }
        end
    }
    nvme_obj.init(obj)
    lu.assertEquals(obj.SerialNumber, 'N/A')
end

function TESTNvmeObject:test_reset_fail_val()
    local obj = {
        failure_debounce = {
            clear_debounced_val = function ()
            end
        },
        pre_failure_debounce = {
            clear_debounced_val = function ()
            end
        }
    }
    nvme_obj.reset_default_fail_val(obj)
    lu.assertEquals(obj.last_fail, 0)
    lu.assertEquals(obj.last_pre_fail, 0)
end

function TESTNvmeObject:test_restore_status_flags()
    local obj = {
        drive_not_ready_debounce = {
            clear_debounced_val = function ()
            end
        },
        drive_functional_debounce = {
            clear_debounced_val = function ()
            end
        },
        port0_pcie_link_active_debounce = {
            clear_debounced_val = function ()
            end
        },
        link_fault_debounce = {
            clear_debounced_val = function ()
            end
        }
    }
    nvme_obj.restore_status_flags(obj)
    lu.assertEquals(obj.drive_not_ready, 1)
    lu.assertEquals(obj.drive_functional, 0)
    lu.assertEquals(obj.reset_not_required, 0)
    lu.assertEquals(obj.port0_pcie_link_active, 0)
end

function TESTNvmeObject:test_get_status_flags()
    local obj = {
        SSDChip = {
            Read = function()
                return '\x06\x02\x03\x04\x05\x06\x07\x3F'
            end
        },
        VPDChip = {},
        drive_not_ready_debounce = {
            get_debounced_val = function()
                return 0
            end
        },
        drive_functional_debounce = {
            get_debounced_val = function()
                return 0
            end
        },
        port0_pcie_link_active_debounce = {
            get_debounced_val = function()
                return 0
            end
        },
        link_fault_debounce = {
            get_debounced_val = function()
                return 0
            end
        }
    }
    nvme_obj.get_status_flags(obj)
    lu.assertEquals(obj.drive_not_ready, 0)
    lu.assertEquals(obj.drive_functional, 0)
    lu.assertEquals(obj.reset_not_required, 0)
    lu.assertEquals(obj.port0_pcie_link_active, 0)
    lu.assertEquals(obj.link_fault, 0)
    lu.assertEquals(obj.Status, 252)
    lu.assertEquals(obj.LifeUsedPercentage, 5)
end

function TESTNvmeObject:test_pciessd_predictive_failure_and_fault_alarm()
    local call_restore_status_flags = false
    local call_get_status_flags = false
    local call_update_link_fault = false
    local obj = {
        restore_status_flags = function()
            call_restore_status_flags = true
        end,
        reset_default_fail_val = function()
        end,
        protocol = common_def.NVME_VPD_PROTOCOL_NVME_MI,
        get_status_flags = function()
            call_get_status_flags = true
        end,
        pciessd_is_enable = function()
            return false
        end,
        update_link_fault = function()
            call_update_link_fault = true
        end,
        get_nvme_pre_failure = function()
        end,
        get_nvme_failure = function()
        end
    }
    local temp_config = c_storageconfig.get_instance
    c_storageconfig.get_instance = function()
        return {
            smbios_status = 3,
            power_state = 'OFF'
        }
    end

    nvme_obj.pciessd_predictive_failure_and_fault_alarm(obj)
    lu.assertEquals(obj.update_time, 0)
    lu.assertEquals(call_restore_status_flags, true)

    c_storageconfig.get_instance = function()
        return {
            smbios_status = 3,
            power_state = 'ON'
        }
    end
    nvme_obj.pciessd_predictive_failure_and_fault_alarm(obj)
    lu.assertEquals(call_get_status_flags, true)

    obj.update_time = 36
    local temp_log = log.notice
    call_log_notice = false
    log.notice = function()
        call_log_notice = true
    end
    nvme_obj.pciessd_predictive_failure_and_fault_alarm(obj)
    lu.assertEquals(call_log_notice, true)

    obj.update_time = 100
    nvme_obj.pciessd_predictive_failure_and_fault_alarm(obj)
    lu.assertEquals(call_update_link_fault, false)

    c_storageconfig.get_instance = temp_config
    log.notice = temp_log
end