-- 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'

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