-- 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 drive_collection = require('drive.drive_collection')
local c_drive = require('drive.drive_object')
local hook_orm_tasks = require('hook_orm_tasks')
local lu = require('luaunit')

TestDriveCollection = {}

function TestDriveCollection:test_init()
    local drive_collection_obj = drive_collection.new()
    local ret = pcall(drive_collection_obj.init, drive_collection_obj)
    lu.assertEquals(ret, true)
end

local function stub_collection_fold(drives)
    local orig = c_drive.collection
    c_drive.collection = {
        fold = function(_, cb, acc)
            for key, obj in pairs(drives) do
                acc = cb(acc, obj, key)
            end
            return acc
        end,
    }
    return orig
end

function TestDriveCollection:test_get_drives_by_controller_id_empty()
    local orig = stub_collection_fold({})
    local dc = drive_collection.new()
    local result = dc:get_drives_by_controller_id(0)
    lu.assertEquals(result, {})
    c_drive.collection = orig
end

function TestDriveCollection:test_get_drives_by_controller_id_filters_correctly()
    local drives = {
        { Name = 'Disk0', RefControllerId = 0, identify_pd = true, presence = 1 },
        { Name = 'Disk1', RefControllerId = 1, identify_pd = true, presence = 1 },
        { Name = 'Disk2', RefControllerId = 0, identify_pd = false, presence = 1 },
        { Name = 'Disk3', RefControllerId = 0, identify_pd = true, presence = 0 },
        { Name = 'Disk4', RefControllerId = 0, identify_pd = true, presence = 1 },
    }
    local orig = stub_collection_fold(drives)
    local dc = drive_collection.new()
    local result = dc:get_drives_by_controller_id(0)
    table.sort(result)
    lu.assertEquals(result, { 'Disk0', 'Disk4' })
    c_drive.collection = orig
end

function TestDriveCollection:test_get_drives_by_controller_id_identify_pd_nil_included()
    local drives = {
        { Name = 'Disk0', RefControllerId = 0, presence = 1 },
    }
    local orig = stub_collection_fold(drives)
    local dc = drive_collection.new()
    local result = dc:get_drives_by_controller_id(0)
    lu.assertEquals(result, { 'Disk0' })
    c_drive.collection = orig
end

function TestDriveCollection:test_get_drives_by_controller_id_no_match()
    local drives = {
        { Name = 'Disk0', RefControllerId = 1, identify_pd = true, presence = 1 },
        { Name = 'Disk1', RefControllerId = 2, identify_pd = true, presence = 1 },
    }
    local orig = stub_collection_fold(drives)
    local dc = drive_collection.new()
    local result = dc:get_drives_by_controller_id(0)
    lu.assertEquals(result, {})
    c_drive.collection = orig
end

function TestDriveCollection:test_get_drives_by_controller_id_returns_independent_tables()
    local drives = {
        { Name = 'Disk0', RefControllerId = 0, identify_pd = true, presence = 1 },
    }
    local orig = stub_collection_fold(drives)
    local dc = drive_collection.new()
    local r1 = dc:get_drives_by_controller_id(0)
    local r2 = dc:get_drives_by_controller_id(0)
    lu.assertNotIs(r1, r2)
    lu.assertEquals(r1, r2)
    c_drive.collection = orig
end

function TestDriveCollection:test_get_drive_objs_by_controller_id_filters_and_returns_objects()
    local drives = {
        { Name = 'Disk0', RefControllerId = 0, identify_pd = true, presence = 1 },
        { Name = 'Disk1', RefControllerId = 1, identify_pd = true, presence = 1 },
        { Name = 'Disk2', RefControllerId = 0, identify_pd = false, presence = 1 },
        { Name = 'Disk3', RefControllerId = 0, identify_pd = true, presence = 0 },
        { Name = 'Disk4', RefControllerId = 0, identify_pd = true, presence = 1 },
    }
    local orig = stub_collection_fold(drives)
    local dc = drive_collection.new()
    local result = dc:get_drive_objs_by_controller_id(0)
    table.sort(result, function(a, b)
        return a.Name < b.Name
    end)
    lu.assertEquals(#result, 2)
    lu.assertEquals(result[1].Name, 'Disk0')
    lu.assertEquals(result[2].Name, 'Disk4')
    c_drive.collection = orig
end

function TestDriveCollection:test_get_drive_objs_by_controller_id_invalid_controller_id_returns_empty()
    local drives = {
        { Name = 'Disk0', RefControllerId = 0, identify_pd = true, presence = 1 },
    }
    local orig = stub_collection_fold(drives)
    local dc = drive_collection.new()
    lu.assertEquals(dc:get_drive_objs_by_controller_id(nil), {})
    lu.assertEquals(dc:get_drive_objs_by_controller_id('0'), {})
    c_drive.collection = orig
end

local function mock_fp_w()
    local lines = {}
    return {
        write = function(_, s)
            lines[#lines + 1] = s
        end,
        content = function(self)
            return table.concat(lines)
        end,
    }
end

local function make_physical_drive_stub(attrs)
    return {
        Id = attrs.Id,
        Presence = attrs.Presence,
        RefControllerId = attrs.RefControllerId,
        dump_physical_drive = function(self, fp_w)
            fp_w:write(string.format('DUMP:%d\n', self.Id))
        end,
    }
end

function TestDriveCollection:test_dump_controller_physical_drive_not_found()
    local orig = stub_collection_fold({})
    local dc = drive_collection.new()
    local fp_w = mock_fp_w()
    dc:dump_controller_physical_drive(0, fp_w)
    lu.assertEquals(fp_w:content(), 'Not Found Physical Drives.\n')
    c_drive.collection = orig
end

function TestDriveCollection:test_dump_controller_physical_drive_none_in_position()
    local drives = {
        make_physical_drive_stub({ Id = 1, Presence = 0, RefControllerId = 0 }),
        make_physical_drive_stub({ Id = 2, Presence = 1, RefControllerId = 1 }),
    }
    local orig = stub_collection_fold(drives)
    local dc = drive_collection.new()
    local fp_w = mock_fp_w()
    dc:dump_controller_physical_drive(0, fp_w)
    lu.assertEquals(fp_w:content(), 'None of physical drives is in position.\n')
    c_drive.collection = orig
end

function TestDriveCollection:test_dump_controller_physical_drive_sorted_by_id()
    local drives = {
        make_physical_drive_stub({ Id = 3, Presence = 1, RefControllerId = 0 }),
        make_physical_drive_stub({ Id = 1, Presence = 1, RefControllerId = 0 }),
        make_physical_drive_stub({ Id = 2, Presence = 1, RefControllerId = 0 }),
        make_physical_drive_stub({ Id = 4, Presence = 1, RefControllerId = 1 }),
        make_physical_drive_stub({ Id = 5, Presence = 0, RefControllerId = 0 }),
    }
    local orig = stub_collection_fold(drives)
    local dc = drive_collection.new()
    local fp_w = mock_fp_w()
    dc:dump_controller_physical_drive(0, fp_w)
    lu.assertEquals(fp_w:content(), 'DUMP:1\nDUMP:2\nDUMP:3\n')
    c_drive.collection = orig
end

function TestDriveCollection:test_dump_controller_physical_drive_single_drive()
    local drives = {
        make_physical_drive_stub({ Id = 7, Presence = 1, RefControllerId = 255 }),
    }
    local orig = stub_collection_fold(drives)
    local dc = drive_collection.new()
    local fp_w = mock_fp_w()
    dc:dump_controller_physical_drive(255, fp_w)
    lu.assertEquals(fp_w:content(), 'DUMP:7\n')
    c_drive.collection = orig
end

-- ============ get_manufacturer：查询物理盘厂商 ============

-- 以 Id 为键构造硬盘对象表，模拟 c_drive.collection 的 find/fold 语义；
-- pcall 包裹保证断言失败时也能恢复 c_drive.collection，不污染后续用例
local function with_stub_collection_find(drives, fn)
    local orig = c_drive.collection
    c_drive.collection = {
        find = function(_, filter)
            return drives[filter.Id]
        end,
        fold = function(_, func, acc)
            for _, drive in pairs(drives) do
                func(acc, drive, nil)
            end
            return acc
        end,
    }
    local ok, err = pcall(fn)
    c_drive.collection = orig
    if not ok then
        error(err)
    end
end

function TestDriveCollection:test_get_manufacturer_found()
    local drives = {
        [1] = { Id = 1, Name = 'Disk1', Manufacturer = 'SSSTC' },
        [2] = { Id = 2, Name = 'Disk2', Manufacturer = 'Intel' },
    }
    with_stub_collection_find(drives, function()
        local dc = drive_collection.new()
        -- DeviceNumber 为 1-based 逻辑序号：按 Id 升序取第 N 个盘
        lu.assertEquals(dc:get_manufacturer(1), 'SSSTC')
        lu.assertEquals(dc:get_manufacturer(2), 'Intel')
    end)
end

function TestDriveCollection:test_get_manufacturer_noncontiguous_id_and_fallback()
    local drives = {
        [1] = { Id = 1, Name = 'Disk1', Manufacturer = 'SSSTC' },
        [2] = { Id = 2, Name = 'Disk2', Manufacturer = 'Intel' },
        [8] = { Id = 8, Name = 'Disk8', Manufacturer = 'Toshiba' },
        [9] = { Id = 9, Name = 'Disk9', Manufacturer = 'WD' },
    }
    with_stub_collection_find(drives, function()
        local dc = drive_collection.new()
        -- 序号 3 → 第 3 个盘 = Id=8（逻辑序号语义，非连续 Id 依然按排序取）
        lu.assertEquals(dc:get_manufacturer(3), 'Toshiba')
        -- 序号 9 越界 → 回退按槽位号 9 直查 → 命中 Id=9
        lu.assertEquals(dc:get_manufacturer(9), 'WD')
    end)
end

function TestDriveCollection:test_get_manufacturer_drive_not_found()
    local drives = {
        [2] = { Id = 2, Name = 'Disk2', Manufacturer = 'Intel' },
    }
    with_stub_collection_find(drives, function()
        local dc = drive_collection.new()
        -- 序号 2 越界（仅 1 个盘）→ 回退槽位 2 → 命中 Id=2
        lu.assertEquals(dc:get_manufacturer(2), 'Intel')
        -- 序号 3 越界 → 回退槽位 3 → 未命中
        local mfr, err_code = dc:get_manufacturer(3)
        lu.assertEquals(mfr, nil)
        lu.assertEquals(err_code, 'not_found')
    end)
end

function TestDriveCollection:test_get_manufacturer_manufacturer_nil()
    local drives = {
        [1] = { Id = 1, Name = 'Disk1' },
    }
    with_stub_collection_find(drives, function()
        local dc = drive_collection.new()
        -- 硬盘对象存在但厂商字段缺失：返回空串
        lu.assertEquals(dc:get_manufacturer(1), '')
    end)
end

function TestDriveCollection:test_get_manufacturer_drive_absent()
    local drives = {
        [1] = { Id = 1, Name = 'Disk1', Manufacturer = 'SSSTC', Presence = 0 },
    }
    with_stub_collection_find(drives, function()
        local dc = drive_collection.new()
        -- 盘不在位：返回 absent 错误码，由 ipmi_service 转 0xCB
        local mfr, err_code = dc:get_manufacturer(1)
        lu.assertEquals(mfr, nil)
        lu.assertEquals(err_code, 'absent')
    end)
end

-- ============ 加固硬盘诊断：诊断队列与串行诊断任务 ============

local function make_diag_drive(name, presence)
    local d = {
        Name = name,
        Presence = presence or 1,
        diagnosed = 0,
    }
    d.diagnose_continue = function(self)
        return self.Presence == 1
    end
    d.diagnose_task = function(self)
        self.diagnosed = self.diagnosed + 1
    end
    return d
end

-- drive_collection.new() 返回共享单例，各用例需先清空队列隔离状态
local function fresh_dc()
    local dc = drive_collection.new()
    dc.diag_queue = {}
    return dc
end

-- 在隔离的 c_tasks 单例中注册并运行诊断周期任务
local function with_hooked_diag_task(dc, fn)
    pcall(hook_orm_tasks.unhook) -- 重置 c_tasks 单例，避免跨用例任务残留
    local my_tasks = hook_orm_tasks.hook()
    dc:start_update_pd_info_task()
    local ok, err = pcall(fn, my_tasks)
    hook_orm_tasks.unhook()
    if not ok then
        error(err)
    end
end

function TestDriveCollection:test_enqueue_fault_drive_fifo_and_dedup()
    local dc = fresh_dc()
    local d1 = make_diag_drive('Disk1')
    dc:enqueue_fault_drive(d1)
    dc:enqueue_fault_drive(d1) -- 重复入队应去重
    dc:enqueue_fault_drive(make_diag_drive('Disk2'))
    lu.assertEquals(#dc.diag_queue, 2)
    lu.assertEquals(dc.diag_queue[1].Name, 'Disk1')
    lu.assertEquals(dc.diag_queue[2].Name, 'Disk2')
end

function TestDriveCollection:test_enqueue_fault_drive_nil_ignored()
    local dc = fresh_dc()
    dc:enqueue_fault_drive(nil)
    dc:enqueue_fault_drive({})
    lu.assertEquals(#dc.diag_queue, 0)
end

function TestDriveCollection:test_reset_diag_queue()
    local dc = fresh_dc()
    dc:enqueue_fault_drive(make_diag_drive('Disk1'))
    dc:reset_diag_queue()
    lu.assertEquals(#dc.diag_queue, 0)
    -- 清空后可重新入队
    dc:enqueue_fault_drive(make_diag_drive('Disk2'))
    lu.assertEquals(#dc.diag_queue, 1)
    -- 空队列清空不报错
    dc:reset_diag_queue()
    lu.assertEquals(#dc.diag_queue, 0)
end

function TestDriveCollection:test_diag_task_serial_fifo()
    local dc = fresh_dc()
    local order = {}
    local d1 = make_diag_drive('Disk1')
    local d2 = make_diag_drive('Disk2')
    d1.diagnose_task = function()
        table.insert(order, 'Disk1')
    end
    d2.diagnose_task = function()
        table.insert(order, 'Disk2')
    end
    dc:enqueue_fault_drive(d1)
    dc:enqueue_fault_drive(d2)
    with_hooked_diag_task(dc, function(my_tasks)
        my_tasks:run_task('drive diagnose task') -- 每轮只出队一个
        lu.assertEquals(order, { 'Disk1' })
        my_tasks:run_task('drive diagnose task')
        lu.assertEquals(order, { 'Disk1', 'Disk2' })
        lu.assertEquals(#dc.diag_queue, 0)
    end)
end

function TestDriveCollection:test_diag_task_skip_not_present()
    local dc = fresh_dc()
    local d1 = make_diag_drive('Disk1', 0) -- 入队后被拔出
    dc:enqueue_fault_drive(d1)
    with_hooked_diag_task(dc, function(my_tasks)
        my_tasks:run_task('drive diagnose task')
        lu.assertEquals(d1.diagnosed, 0)
        lu.assertEquals(#dc.diag_queue, 0)
    end)
end

function TestDriveCollection:test_diag_task_error_not_blocking_next()
    local dc = fresh_dc()
    local d1 = make_diag_drive('Disk1')
    d1.diagnose_task = function()
        error('diag fail')
    end
    local d2 = make_diag_drive('Disk2')
    dc:enqueue_fault_drive(d1)
    dc:enqueue_fault_drive(d2)
    with_hooked_diag_task(dc, function(my_tasks)
        my_tasks:run_task('drive diagnose task') -- d1 异常被 pcall 捕获
        my_tasks:run_task('drive diagnose task') -- d2 正常执行
        lu.assertEquals(d2.diagnosed, 1)
    end)
end

function TestDriveCollection:test_smbios_status_changed_clears_diag_queue()
    local orig = stub_collection_fold({})
    local dc = fresh_dc()
    dc:process_smbios_status_changed()
    dc:enqueue_fault_drive(make_diag_drive('Disk1'))
    dc.on_smbios_status_changed:emit(false) -- 下电清空队列
    lu.assertEquals(#dc.diag_queue, 0)
    dc:enqueue_fault_drive(make_diag_drive('Disk2'))
    dc.on_smbios_status_changed:emit(true) -- 上电清空队列
    lu.assertEquals(#dc.diag_queue, 0)
    c_drive.collection = orig
end