local lu = require('luaunit')
local common_def = require('common_def')
local log = require('mc.logging')
local nvme_obj = require('nvme.nvme_object')

TestHandlerNvme = {}

function TestHandlerNvme:setUp()
    self.handler = require('bma.handles.handler_nvme')
    self.log_records = {}
    self.orig_maintenance = log.maintenance
    log.maintenance = function(_, level, fc, fmt, ...)
        self.log_records[#self.log_records + 1] = string.format(fmt, ...)
    end
end

function TestHandlerNvme:tearDown()
    log.maintenance = self.orig_maintenance
end

local function create_mock_nvme(overrides)
    local obj = {
        SerialNumber = 'TEST-SN-001',
        Slot = 3,
        CapacityMiB = common_def.INVALID_U32,
        NegotiatedSpeedGbs = common_def.INVALID_U8,
        PowerOnHours = common_def.INVALID_U32,
        MediaErrorCount = common_def.INVALID_U32,
        SpareBlockPercentage = common_def.INVALID_U8,
        CriticalWarning = common_def.INVALID_U8,
        Revision = 'N/A',
        t_prop_without_bma = {},
        capacitor_fault_active = false,
        capacitor_fault_last_log_time = 0,
        spare_block_fault_active = false,
        spare_block_fault_last_log_time = 0,
        fused_smart_info = {
            percent_used = {
                HWSmart = common_def.INVALID_U8,
                Smart = common_def.INVALID_U8,
                BMA = common_def.INVALID_U8,
            },
            host_write_l = {
                HWSmart = common_def.STORAGE_INFO_INVALID_DWORD,
                Smart = common_def.STORAGE_INFO_INVALID_DWORD,
                BMA = common_def.STORAGE_INFO_INVALID_DWORD,
            },
            host_write_h = {
                HWSmart = common_def.STORAGE_INFO_INVALID_DWORD,
                Smart = common_def.STORAGE_INFO_INVALID_DWORD,
                BMA = common_def.STORAGE_INFO_INVALID_DWORD,
            },
            nand_write_l = {
                HWSmart = common_def.STORAGE_INFO_INVALID_DWORD,
                Smart = common_def.STORAGE_INFO_INVALID_DWORD,
                BMA = common_def.STORAGE_INFO_INVALID_DWORD,
            },
            nand_write_h = {
                HWSmart = common_def.STORAGE_INFO_INVALID_DWORD,
                Smart = common_def.STORAGE_INFO_INVALID_DWORD,
                BMA = common_def.STORAGE_INFO_INVALID_DWORD,
            },
        },
    }
    setmetatable(obj, { __index = nvme_obj })
    if overrides then
        for k, v in pairs(overrides) do
            obj[k] = v
        end
    end
    return obj
end

function TestHandlerNvme:test_update_capacity_mb()
    local obj = create_mock_nvme()
    local data = { CapacityBytes = 1024 * 1024 * 500 }
    self.handler:update_capacity_mb(obj, data)
    lu.assertEquals(obj.CapacityMiB, 500)
end

function TestHandlerNvme:test_update_capacity_mb_nil()
    local obj = create_mock_nvme()
    self.handler:update_capacity_mb(obj, {})
    lu.assertEquals(obj.CapacityMiB, common_def.INVALID_U32)
end

function TestHandlerNvme:test_update_speed_gb()
    local obj = create_mock_nvme()
    local data = { NegotiatedSpeedGbs = 8.0 }
    self.handler:update_speed_gb(obj, data)
    lu.assertEquals(obj.NegotiatedSpeedGbs, 7)
end

function TestHandlerNvme:test_update_speed_gb_nil()
    local obj = create_mock_nvme()
    self.handler:update_speed_gb(obj, {})
    lu.assertEquals(obj.NegotiatedSpeedGbs, common_def.INVALID_U8)
end

function TestHandlerNvme:test_update_smart_info_basic()
    local obj = create_mock_nvme()
    local data = {
        SMARTInfo = {
            PowerOnHours = 1000,
            MediaErrorCount = 5,
            AvailableSpare = 80,
            CriticalWarning = 0,
            AvailableSpareThreshold = 10,
        },
    }
    self.handler:update_smart_info(obj, data)
    lu.assertEquals(obj.PowerOnHours, 1000)
    lu.assertEquals(obj.MediaErrorCount, 5)
    lu.assertEquals(obj.SpareBlockPercentage, 80)
    lu.assertEquals(obj.CriticalWarning, 0)
end

function TestHandlerNvme:test_update_smart_info_nil_smart()
    local obj = create_mock_nvme()
    self.handler:update_smart_info(obj, {})
    lu.assertEquals(obj.PowerOnHours, common_def.INVALID_U32)
    lu.assertEquals(obj.MediaErrorCount, common_def.INVALID_U32)
end

function TestHandlerNvme:test_update_smart_info_capacitor_fault_triggered()
    local obj = create_mock_nvme()
    local data = {
        SMARTInfo = {
            PowerOnHours = 1000,
            MediaErrorCount = 5,
            AvailableSpare = 80,
            CriticalWarning = 0x10,
            AvailableSpareThreshold = 10,
        },
    }
    self.handler:update_smart_info(obj, data)
    lu.assertEquals(obj.CriticalWarning, 0x10)
end

function TestHandlerNvme:test_update_smart_info_spare_block_log()
    local obj = create_mock_nvme()
    local data = {
        SMARTInfo = {
            PowerOnHours = 1000,
            MediaErrorCount = 5,
            AvailableSpare = 5,
            CriticalWarning = 0,
            AvailableSpareThreshold = 10,
        },
    }
    self.handler:update_smart_info(obj, data)
    lu.assertEquals(obj.SpareBlockPercentage, 5)
    lu.assertEquals(obj.spare_block_fault_active, true)
    lu.assertTrue(#self.log_records > 0)
end

function TestHandlerNvme:test_reset_nvme()
    local obj = create_mock_nvme({
        PowerOnHours = 1000,
        MediaErrorCount = 5,
        NegotiatedSpeedGbs = 7,
        SpareBlockPercentage = 80,
        CriticalWarning = 0x10,
        spare_block_fault_active = true,
        spare_block_fault_last_log_time = 1000,
        capacitor_fault_active = true,
        capacitor_fault_last_log_time = 1000,
    })
    self.handler:reset_nvme(obj)
    lu.assertEquals(obj.PowerOnHours, common_def.INVALID_U32)
    lu.assertEquals(obj.MediaErrorCount, common_def.INVALID_U32)
    lu.assertEquals(obj.NegotiatedSpeedGbs, common_def.INVALID_U8)
    lu.assertEquals(obj.SpareBlockPercentage, common_def.INVALID_U8)
end

function TestHandlerNvme:test_update_nvme_full()
    local obj = create_mock_nvme()
    local data = {
        CapacityBytes = 1024 * 1024 * 200,
        NegotiatedSpeedGbs = 16.0,
        SMARTInfo = {
            PowerOnHours = 500,
            MediaErrorCount = 2,
            AvailableSpare = 90,
            CriticalWarning = 0,
            AvailableSpareThreshold = 10,
        },
        FirmwareVersion = '1.0.0',
    }
    self.handler:update_nvme(obj, data)
    lu.assertEquals(obj.CapacityMiB, 200)
    lu.assertEquals(obj.NegotiatedSpeedGbs, 9)
    lu.assertEquals(obj.PowerOnHours, 500)
    lu.assertEquals(obj.MediaErrorCount, 2)
    lu.assertEquals(obj.SpareBlockPercentage, 90)
end

function TestHandlerNvme:test_add_calls_update_nvme()
    local obj = create_mock_nvme()
    local data = {
        CapacityBytes = 1024 * 1024 * 100,
        NegotiatedSpeedGbs = 5.0,
        SMARTInfo = {
            PowerOnHours = 200,
            MediaErrorCount = 1,
            AvailableSpare = 95,
            CriticalWarning = 0,
            AvailableSpareThreshold = 10,
        },
        FirmwareVersion = '2.0.0',
    }
    self.handler:add(nil, data, obj)
    lu.assertEquals(obj.CapacityMiB, 100)
    lu.assertEquals(obj.PowerOnHours, 200)
end

function TestHandlerNvme:test_update_smart_info_fused_bma()
    local obj = create_mock_nvme()
    local data = {
        SMARTInfo = {
            PowerOnHours = 1000,
            MediaErrorCount = 5,
            AvailableSpare = 80,
            CriticalWarning = 0,
            AvailableSpareThreshold = 10,
            PercentageUsed = 30,
            DataUnitsWritten = 1024.5,
            PeriodWriteCount = 2048.7,
        },
    }
    self.handler:update_smart_info(obj, data)
    lu.assertEquals(obj.fused_smart_info.percent_used.BMA, 30)
    lu.assertEquals(obj.fused_smart_info.host_write_l.BMA, math.floor(1024.5 * 1024 / 512))
    lu.assertEquals(obj.fused_smart_info.host_write_h.BMA, 0)
    lu.assertEquals(obj.fused_smart_info.nand_write_l.BMA, math.floor(2048.7 * 1024 / 512))
    lu.assertEquals(obj.fused_smart_info.nand_write_h.BMA, 0)
end

function TestHandlerNvme:test_update_smart_info_fused_bma_nil()
    local obj = create_mock_nvme()
    local data = {
        SMARTInfo = {
            PowerOnHours = 1000,
            MediaErrorCount = 5,
        },
    }
    self.handler:update_smart_info(obj, data)
    lu.assertEquals(obj.fused_smart_info.percent_used.BMA, common_def.INVALID_U8)
    lu.assertEquals(obj.fused_smart_info.host_write_l.BMA, common_def.STORAGE_INFO_INVALID_DWORD)
    lu.assertEquals(obj.fused_smart_info.nand_write_l.BMA, common_def.STORAGE_INFO_INVALID_DWORD)
end

function TestHandlerNvme:test_update_calls_update_nvme()
    local obj = create_mock_nvme()
    local data = {
        CapacityBytes = 1024 * 1024 * 300,
        NegotiatedSpeedGbs = nil,
        SMARTInfo = nil,
        FirmwareVersion = nil,
    }
    self.handler:update(nil, data, obj)
    lu.assertEquals(obj.CapacityMiB, 300)
end

function TestHandlerNvme:test_delete_calls_reset_nvme()
    local obj = create_mock_nvme({
        PowerOnHours = 1000,
        MediaErrorCount = 5,
    })
    self.handler:delete(nil, nil, obj)
    lu.assertEquals(obj.PowerOnHours, common_def.INVALID_U32)
    lu.assertEquals(obj.MediaErrorCount, common_def.INVALID_U32)
end

function TestHandlerNvme:test_reset_calls_reset_nvme_for_all_objects()
    local h = self.handler
    local orig_objects = h.objects
    local mock_obj1 = create_mock_nvme({ PowerOnHours = 100 })
    local mock_obj2 = create_mock_nvme({ PowerOnHours = 200 })
    h.objects = {
        [mock_obj1] = '/path1',
        [mock_obj2] = '/path2',
    }
    h:reset()
    lu.assertEquals(mock_obj1.PowerOnHours, common_def.INVALID_U32)
    lu.assertEquals(mock_obj2.PowerOnHours, common_def.INVALID_U32)
    h.objects = orig_objects
end