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

-- Description: 升级信号的回调处理
local class = require 'mc.class'
local singleton = require 'mc.singleton'
local fw_mgmt = require 'power_mgmt.client'
local ctx = require 'mc.context'
local psu_service = require 'psu_service'
local c_psu_object = require 'device.psu'
local c_power_configuration = require 'device.power_configuration'
local fw_def = require 'macros.fw_def'
local c_tasks = require 'mc.orm.tasks'
local log = require 'mc.logging'
local signal = require 'signal'
local client = require 'power_mgmt.client'
local context = require 'mc.context'
local utils = require 'power_mgmt_utils'
local lock_chip = require 'lock_chip'
local enums = require 'macros.power_mgmt_enums'
local c_power_supplies = require 'device.power_supplies'

local RET_ERR<const> = -1
local RET_UPGRADE_EXIST<const> = 36
local E_FAILED <const> = '' -- 空错误信息
local UPGRADE_PERCENT_INIT<const> = 8
local UPGRADE_PERCENT_AFTER_LOAD <const> = 98 -- 电源升级完成的总进度设置为98%,防止在100%时长时间等待
local COMP_POWER_TYPE <const> = 'Power'

local NOT_IN_UPGRADE_STATUS = 0
local PREPARE_UPGRADE_STATUS = 1
local IN_UPGRADE_STATUS = 2

local COMPONENT_IDEX<const> = {
    ['pmbus'] = 0,
    ['canbus'] = 2
}

-- upgrade_process 电源固件升级进度回调函数
local upgrade_process = {
    step_progress = 0,
    cur_progress = 0,
    start_process = 0,
    init = function(self, init_progress)
        self.cur_progress = init_progress
        self.step_progress = init_progress
        self.start_process = init_progress
    end,
    get_comp_percent = function(self)
        return self.cur_progress
    end,
    set_comp_percent = function(self, progress, para_tab)
        progress = math.floor(progress)
        if progress > 100 then
            progress = 100
        end
        if self.cur_progress == progress then
            return
        end
        self.cur_progress = progress
        fw_mgmt:UpdateServiceUpdateServiceUpdateUpgradeStatus(ctx.new(), 1, 'Power', 0, progress, 'Upgrading', para_tab)
    end,
    set_step_progress = function(self, progress)
        -- 升级进度不超过100
        if progress > 100 then
            progress = 100
        end
        self.start_process = self.step_progress
        self.step_progress = progress
    end,
    get_step_progress = function(self)
        return self.step_progress
    end
}

local function extract_package(file_path)
    local path_split = #file_path - file_path:reverse():find('/') + 1 -- 执行失败，认为是升级失败，在上层pcall处返回
    local upgrade_path = file_path:sub(1, path_split)
    local upgrade_ps_packet = file_path:sub(path_split + 1, -1)
    local ok = utils.secure_tar_unzip(
        upgrade_path .. upgrade_ps_packet,
        upgrade_path,
        0x6400000, -- 最大解压限制100M
        1024
    ) -- 解压tar包
    if not ok then
        log:error('Secure unzip failed.')
        return nil
    end
    return upgrade_path
end

-- 获取是否支持单电源升级
local function is_singal_psu_support()
    local power_configuration = c_power_configuration.collection:find({})
    if not power_configuration then
        return false
    end
    return power_configuration:get_single_ps_support_flag() == 1
end

-- 更新电源升级状态
local function update_upgrade_status(upgrade_status)
    local power_supplies = c_power_supplies.collection:find({})
    if power_supplies then
        log:notice("update_upgrade_status %s", upgrade_status)
        power_supplies:set_upgrade_status(upgrade_status)
    end
end

local function is_support_upgrade(psu_obj)
    -- 单电源环境支持升级 或者 整机功耗支持升级
    return is_singal_psu_support() or psu_service.get_instance():check_ps_num_exclude_ps({[psu_obj.ps_id] = 1})
end

local function upgrade_psu(psu_obj, upgrade_path, process, para_tab, upgrade_channel)
    if psu_obj.LossOfInput ~= 0 then
        log:error('[power_mgmt] upgrade ps%s failed, input loss state is %s', psu_obj.ps_id, psu_obj.LossOfInput)
        return fw_def.SUPPLY_3
    end
    local chip = psu_obj.psu_slot.PsuChip[enums.PSU_CHIP_OBJECT.CHIP]
    if upgrade_channel == 1 then
        chip = psu_obj.psu_slot.PsuBackupChip[enums.PSU_CHIP_OBJECT.CHIP]
    end
    local ok, res = pcall(function ()
        return lock_chip.get_instance():retry_lock_chip(chip)
    end)
    if not ok then
        log:error('[power_mgmt] upgrade ps%d failed, ret = %s', psu_obj.ps_id, res)
        return false
    end
    local ret
    for retry_time = 1, fw_def.UPGRADE_ONE_PSU_RETRY_TIMES do
        ret = psu_obj:power_upgrade(upgrade_path, process, para_tab)
        if ret ~= fw_def.SUPPLY_1 then
            break
        end
        log:notice("Ps%s upgrade failed, upgrade times is %s.", psu_obj.ps_id, retry_time)
    end
    ok, res = pcall(function ()
        return lock_chip.get_instance():retry_unlock_chip(chip)
    end)
    -- 认为升级执行完成，仅解锁失败，记录解锁失败日志
    if not ok then
        log:error('[power_mgmt] upgrade ps%d failed, ret = %s', psu_obj.ps_id, res)
    end
    return ret
end

local function upgrade_one_psu(upgrade_path, psu_obj, pro_percent, para_tab, upgrade_channel)
    local upgrade_result = fw_def.ERROR_CODE.CommonFinish
    upgrade_process:set_step_progress(upgrade_process:get_step_progress() + pro_percent)
    log:notice('[power_mgmt] upgrade ps%d', psu_obj.ps_id)

    if not is_support_upgrade(psu_obj) then
        log:error('[power_mgmt] power supply is less than required and not support single power supply')
        upgrade_result = fw_def.ERROR_CODE.InsufficientPowerSupply
        return false, upgrade_result
    end
    local ret = upgrade_psu(psu_obj, upgrade_path, upgrade_process, para_tab, upgrade_channel)
    if ret ~= fw_def.SUPPLY_0 then
        if ret ~= fw_def.SUPPLY_3 and ret ~= fw_def.SUPPLY_6 then
            log:error('[power_mgmt] upgrade ps%d failed, ret = %s', psu_obj.ps_id, ret)
            upgrade_result = fw_def.ERROR_CODE.FirmwareUpgradeError
            return false, upgrade_result
        end
        log:warn('[power_mgmt] ps%d upgrading unsupported, ret = %s', psu_obj.ps_id, ret)
    else
        log:notice('[power_mgmt] ps%d upgrade successfully', psu_obj.ps_id)
    end
    return true, upgrade_result, ret
end

local function check_upgrade_result(result, psu_num)
    local supply_6_num = 0
    local supply_0_num = 0

    -- 一个机型可能有多种型号电源，遍历电源升级时，会出现电源升级包和电源型号不匹配的场景，即SUPPLY_6，不属于升级失败
    -- 需要分别统计SUPPLY_0（升级成功）和 SUPPLY_6（电源升级hpm包和电源型号不匹配）的电源个数
    for psu_id, psu_res in pairs(result) do
        if psu_res == fw_def.SUPPLY_6 then
            supply_6_num = supply_6_num + 1
        elseif psu_res == fw_def.SUPPLY_0 then
            supply_0_num = supply_0_num + 1
        end
    end

    log:notice('check_upgrade_result supply_6_num %s supply_0_num %s', supply_6_num, supply_0_num)

    -- 所有电源升级报错都是SUPPLY_6（电源升级hpm包和电源型号不匹配）， 没有电源升级成功，此升级包是无效的升级包，返回升级失败
    if supply_6_num == psu_num then
        return fw_def.ERROR_CODE.FirmwareUpgradeError
    -- 所有电源都升级成功，返回升级成功
    elseif supply_0_num == psu_num then
        return fw_def.ERROR_CODE.CommonFinish
    -- 部分电源升级报错是SUPPLY_6（电源升级hpm包和电源型号不匹配），其余都升级成功
    -- 电源升级包中类型的电源都已升级成功，返回升级成功
    elseif supply_6_num + supply_0_num == psu_num then
        return fw_def.ERROR_CODE.CommonFinish
    -- 如果有其他类型升级报错，返回升级失败
    else
        return fw_def.ERROR_CODE.FirmwareUpgradeError
    end
end

--支持双通道的电源需要再升级前设置其余电源的轮询通道，不支持则直接返回
local function upgrade_set_psu_channle(psu_obj)
    local upgrade_channel = psu_obj:get_upgrade_channel()
    if upgrade_channel == E_FAILED then --返回E_FAILED说明不支持双通道功能，不用切换备用通道
        return E_FAILED
    end
    c_psu_object.collection:fold(function (_, psus_obj)
        if psus_obj == psu_obj then
            return --升级板子无须切换通道
        end
        psus_obj:change_polling_channel(upgrade_channel)
    end)
    return upgrade_channel
end

--升级完恢复电源轮询通道
local function upgrade_recover_psu_channle(psu_obj, upgrade_channel)
    if upgrade_channel == E_FAILED then
        return nil --upgrade_channel 为 E_FAILED说明不支持双通道功能，不用切换备用通道
    end
    c_psu_object.collection:fold(function (_, psus_obj)
        if psus_obj == psu_obj then
            return --升级板子无须恢复通道
        end
        psus_obj:recover_polling_channel()
    end)
    return nil -- 函数执行成功返回nil
end

-- 依次升级每个psu
local function upgrade_each_psu(file_path, psu_objs, pro_percent, para_tab)
    local upgrade_result = fw_def.ERROR_CODE.CommonFinish
    local upgrade_path = extract_package(file_path)
    local result = {}
    if not upgrade_path then
        utils.get_instance():set_upgrade_flag(false)
        upgrade_result = fw_def.ERROR_CODE.InvalidFirmwarePackage
        return upgrade_result
    end
    local res, res1, ret
    for _, psu_obj in ipairs(psu_objs) do
        res1 = upgrade_set_psu_channle(psu_obj)
        res, upgrade_result, ret = upgrade_one_psu(upgrade_path, psu_obj, pro_percent, para_tab, res1)
        upgrade_recover_psu_channle(psu_obj, res1)
        if res == false then
            utils.get_instance():set_upgrade_flag(false)
            utils.operation_log(nil, 'Upgrade psu(%s) failed', psu_obj.ps_id)
            break
        end
        result[psu_obj.ps_id] = ret
    end

    -- 根据每个电源的升级结果，对最终升级结果进行校验
    if upgrade_result == fw_def.ERROR_CODE.CommonFinish then
        upgrade_result = check_upgrade_result(result, #psu_objs)
        log:notice('[power_mgmt] upgrade result %s', upgrade_result)
    end
    return upgrade_result
end

local function get_upgrade_psu(psu_objs)
    local psu_ids = utils.get_instance():get_power_upgrade_slot_number()
    if not next(psu_ids) then
        return psu_objs
    end
    local psu_id_map = {}
    for _, psu_obj in pairs(psu_objs) do
        psu_id_map[psu_obj.ps_id] = psu_obj
    end
    local res = {}
    local nonexistent_ps_ids = {}
    for _, id in pairs(psu_ids) do
        if not psu_id_map[id] then
            table.insert(nonexistent_ps_ids, id)
        else
            table.insert(res, psu_id_map[id])
        end    
    end
    if next(nonexistent_ps_ids) then
        log:error('The specified power slot [%s] does not exist, please check the input parameters',
            table.concat(nonexistent_ps_ids, ' '))
        utils.operation_log(nil, 'Upgrade (Power) failed, psu(%s) does not exist',
            table.concat(nonexistent_ps_ids, ' '))
        return {}
    end
    return res
end

local function upgrade_power_process(system_id, firmware_type, file_path, component_idex, para_tab)
    -- 计算电源升级进度
    local psu_objs = c_psu_object.collection:fetch(function (object)
        -- 获取协议和升级包匹配且支持升级的电源
        return COMPONENT_IDEX[object.Protocol] == component_idex and object.UpgradeSupported
    end)

    local count = #psu_objs
    log:notice('upgrade_process psu count %d', count)
    local upgrade_result
    if count == 0 then
        log:error('no psu need upgrade')
        utils.get_instance():set_upgrade_flag(false)
        upgrade_result = fw_def.ERROR_CODE.FirmwareUpgradeError
        return upgrade_result
    end
    if utils.get_instance():get_perform_upgrade_flag() == false then
        log:error('The input parameter is wrong. It must be in the following format. Single power supply: slot.' ..
            'Multi-power supply: slot,slot,slot. And slot must be unique and a natural number.')
        upgrade_result = fw_def.ERROR_CODE.FirmwareUpgradeError
        utils.get_instance():set_upgrade_flag(false)
        utils.operation_log(nil, 'Upgrade (Power) failed, please check the input format of slot number')
        return upgrade_result
    end
    local pro_percent = (UPGRADE_PERCENT_AFTER_LOAD - upgrade_process:get_comp_percent()) // count
    -- 升级过程停止数据监测
    psu_service.get_instance().upgrade_flag = 1
    c_psu_object.collection:fold(function (_, obj)
        -- 升级切换到主用，防止掉电；升级结束，能效会切主备
        obj:set_work_mode(0)
    end)
    local upgrade_psu_objs = get_upgrade_psu(psu_objs)
    if not next(upgrade_psu_objs) then
        utils.get_instance():set_upgrade_flag(false)
        upgrade_result = fw_def.ERROR_CODE.FirmwareUpgradeError
        return upgrade_result
    end
    utils.get_instance():set_upgrade_psu_objs(upgrade_psu_objs)
    for _, psu_obj in pairs(upgrade_psu_objs) do
        psu_obj.IsUpgrading = true
    end

    c_tasks.sleep_ms(3000) -- 等待任务队列清空,电压拉偏完成

    -- 优先升级不健康的电源，保证供电
    table.sort(
        upgrade_psu_objs,
        function (obj1, obj2)
            return not obj1:is_healthy() and obj2:is_healthy()
        end
    )
    upgrade_result = upgrade_each_psu(file_path, upgrade_psu_objs, pro_percent, para_tab)
    psu_service.get_instance().upgrade_flag = 0
    for _, psu_obj in pairs(upgrade_psu_objs) do
        psu_obj.IsUpgrading = false
    end
    return upgrade_result
end

local function prepare_upgrade(system_id, firmware_type, cfg_path, file_path, para_tab)
    if firmware_type ~= COMP_POWER_TYPE then
        return
    end
    local ok, rsp = client:PFileFileChown(context.new(), nil, cfg_path, 104, 104)  --设置配置文件属主为secbox
    if not ok then
        log:error('[power_mgmt] chown cfg file failed, error %s', rsp)
        client:UpdateServiceUpdateServicePrepareReply(context.new(), system_id, firmware_type, '', RET_ERR, para_tab)
        return
    end

    -- 检查ex_id
    utils.cfgs:get_cfgs(cfg_path)
    if utils.cfgs:check_idex(0x7D) then
        return
    end
    update_upgrade_status(PREPARE_UPGRADE_STATUS)
    local firmware_version = utils.cfgs:get_version()
    if not firmware_version then
        log:error('Get firmware version failed')
        return
    end
    utils.get_instance():set_update_cfg_firmware_version(firmware_version)
    if utils.get_instance():get_upgrade_flag() then
        log:error('[power_mgmt] power is in upgrading process')
        client:UpdateServiceUpdateServicePrepareReply(context.new(), system_id, firmware_type, '',
            RET_UPGRADE_EXIST, para_tab)
    end
    utils.get_instance():set_upgrade_flag(true)
    log:notice('power upgrade prepare')
    -- prepare阶段升级进度为8%
    upgrade_process:init(UPGRADE_PERCENT_INIT)
    -- 在prepare阶段初始化指定电源升级标志为true，表示升级所有电源
    utils.get_instance():set_perform_upgrade_flag(true)
    -- 在prepare阶段初始化指定电源的升级槽位号，即清除上一次指定电源的升级槽位
    utils.get_instance():init_power_upgrade_slot_number()
    -- 在prepare阶段初始化需要升级的电源对象为nil
    utils.get_instance():set_upgrade_psu_objs(nil)
    if para_tab.SlotNumber then
        -- 在prepare阶段根据传入的SlotNumber确认是否可以执行指定电源升级
        utils.get_instance():get_power_upgrade_flag_and_slot_number(para_tab.SlotNumber)
    end
    fw_mgmt:UpdateServiceUpdateServicePrepareReply(ctx.new(), system_id, firmware_type, '', 0, para_tab)
end

local function process_upgrade(db, system_id, firmware_type, file_path, para_tab)
    update_upgrade_status(IN_UPGRADE_STATUS)
    if firmware_type ~= COMP_POWER_TYPE then
        update_upgrade_status(NOT_IN_UPGRADE_STATUS)
        return
    end
    if utils.cfgs:check_idex(0x7D) then
        update_upgrade_status(NOT_IN_UPGRADE_STATUS)
        return
    end
    local component_idex = utils.cfgs:get_component_idex()
    local ok, rsp = client:PFileFileChown(context.new(), nil, file_path, 104, 104)  --设置配置文件属主为secbox
    if not ok then
        update_upgrade_status(NOT_IN_UPGRADE_STATUS)
        log:error('[power_mgmt] chown file failed, error %s', rsp)
        utils.get_instance():set_upgrade_flag(false)
        client:PUpdateServiceUpdateServiceProcessReply(context.new(), system_id, firmware_type, RET_ERR, para_tab)
        return
    end
    log:notice('power upgrade process')
    local ok, ret = pcall(upgrade_power_process, system_id, firmware_type, file_path, component_idex, para_tab)
    log:notice('[power_mgmt] power upgrade, ret: %s', ret)
    if ret == fw_def.ERROR_CODE.CommonFinish then
        utils.get_instance():set_power_upgrade_success_flag(true)
    else
        utils.get_instance():set_power_upgrade_success_flag(false)
    end
    ret = ok and ret or fw_def.ERROR_CODE.FirmwareUpgradeError
    fw_mgmt:UpdateServiceUpdateServiceProcessReply(ctx.new(), system_id, firmware_type, ret, para_tab)
    update_upgrade_status(NOT_IN_UPGRADE_STATUS)
end

local function finish_upgrade(system_id, firmware_type, para_tab)
    if firmware_type ~= COMP_POWER_TYPE then
        return
    end
    if utils.cfgs:check_idex(0x7D) then
        return
    end
    utils.get_instance():set_upgrade_flag(false)
    log:notice('power upgrade finish')
    local psu_ids = utils.get_instance():get_power_upgrade_slot_number()
    -- power升级完成后，对固件版本号进行校验，校验成功则返回0，否则返回失败
    local firmware_version = utils.get_instance():get_update_cfg_firmware_version()
    if firmware_version == 'null' then
        goto continue
    end
    for _, psu_obj in pairs(utils.get_instance():get_upgrade_psu_objs()) do
        if psu_obj.FirmwareVersion ~= firmware_version then
            log:error('Failed to verify the firmware version, the expected firmware version is %s, '..
                'but psu[%s] actual value is %s', firmware_version, psu_obj.ps_id, psu_obj.FirmwareVersion)
            utils.operation_log(nil, 'Upgrade psu(%s) failed', next(psu_ids) and table.concat(psu_ids, ' ') or 'all')
            fw_mgmt:UpdateServiceUpdateServiceFinishReply(ctx.new(), system_id, firmware_type,
                fw_def.ERROR_CODE.FirmwareFileMismatch, para_tab)
            return
        end
    end
    ::continue::
    if utils.get_instance():get_power_upgrade_success_flag() then
        utils.operation_log(nil, 'Upgrade psu(%s) successfully', next(psu_ids) and table.concat(psu_ids, ' ') or 'all')
    end
    fw_mgmt:UpdateServiceUpdateServiceFinishReply(ctx.new(), system_id, firmware_type, 0, para_tab)
end

local power_upgrade = class()
power_upgrade.prepare_upgrade = prepare_upgrade
power_upgrade.process_upgrade = process_upgrade
power_upgrade.finish_upgrade = finish_upgrade
power_upgrade.upgrade_power_process = upgrade_power_process
power_upgrade.check_upgrade_result = check_upgrade_result
power_upgrade.upgrade_each_psu = upgrade_each_psu
power_upgrade.upgrade_set_psu_channle = upgrade_set_psu_channle
power_upgrade.upgrade_recover_psu_channle = upgrade_recover_psu_channle
power_upgrade.upgrade_process = upgrade_process
power_upgrade.upgrade_psu = upgrade_psu

function power_upgrade:ctor()
    signal.register_callback('upgrade_prepare_callback', nil, prepare_upgrade)
    signal.register_callback('upgrade_process_callback', nil, process_upgrade)
    signal.register_callback('upgrade_finish_callback', nil, finish_upgrade)
end

return singleton(power_upgrade)
