-- 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 class = require 'mc.class'
local c_gpu = require 'gpu_service.gpu_object'
local skynet = require 'skynet'
local log = require 'mc.logging'
local context = require 'mc.context'
local utils = require 'mc.utils'
local bs = require 'mc.bitstring'
 
---@class VastaiGPUCard: GPU @瀚博GPU卡
local c_vastai_gpu = class(c_gpu)
 
local GET_SN_OFFSET <const> = 0x01
local GET_SN_LEN <const> = 0x11
local GET_POWER_OFFSET <const> = 0x30
local GET_POWER_LEN <const> = 0x05
local SENSOR_NA_READING <const> = 0x4000            -- 不在位或者读不到值时，属性设为0x4000
local SENSOR_INVALID_READING <const> = 0x8000       -- 读取失败时，属性设为0x8000
 
local function retry_chip_read(wait, retries, chip, cmd, len)
    local ok = false
    local rsp
    for _ = 1, retries do
        ok, rsp = pcall(function ()
            return chip:Read(context.get_context_or_default(), cmd, len)
        end)
        if ok then
            return ok, rsp
        end
        skynet.sleep(wait)
    end
    return ok, rsp
end
 
local function vastai_verify_data(data)
    local sum = 0
    for i = 1, #data do
        sum = sum + string.sub(data, i, i):byte()
    end
    -- 第一位为校验和;对接收到的数据(包括校验和)进行累加之后再加1,如果得到0,则说明数据未出现传输错误
    if (sum + 1) % 0x0100 == 0 then
        return true
    end
    return false
end
 
function c_vastai_gpu:get_vastai_sn(chip)
    local ok, data = retry_chip_read(10, 10, chip, GET_SN_OFFSET, GET_SN_LEN)
    if not ok then
        log:error('[VastaiGPU]get SN failed, error: %s', data)
        return ""
    end
    if not vastai_verify_data(data) then
        log:error('[VastaiGPU]sn verify data failed, data is %s', utils.to_hex(data))
        return ""
    end
    -- 第一位为校验位,SN从第二位开始
    return string.sub(data, 2)
end
 
function c_vastai_gpu:get_vastai_power(chip)
    local ok, data = retry_chip_read(10, 10, chip, GET_POWER_OFFSET, GET_POWER_LEN)
    if not ok then
        log:error('[VastaiGPU]get power failed, error: %s', data)
        return 0
    end
    if not vastai_verify_data(data) then
        log:error('[VastaiGPU]pwoer verify data failed, data is %s', utils.to_hex(data))
        return 0
    end
    local power_data = bs.new([[<<verify:8, power:32>>]]):unpack(data, true).power
    -- 功耗为有效值时需从uW转换成W
    if power_data ~= SENSOR_NA_READING and power_data ~= SENSOR_INVALID_READING then
        power_data = power_data // 1000000
    end
    return power_data
end
 
function c_vastai_gpu:stop_action_task()
    if not self.action_task then
        return
    end
    self.action_task = false
end
 
function c_vastai_gpu:change_vastai_info()
    local sn, power
    self.init_fork_vastai_info = true
    while true do
        if self.action_task then
            sn = self:get_vastai_sn(self.gpu.RefChip)
            if sn and sn ~= "" then
                self:set_prop('SN', sn)
            end
            power = self:get_vastai_power(self.gpu.RefChip)
            self:set_prop('PowerWatts', power)
        end
        skynet.sleep(1000)
    end
end
 
function c_vastai_gpu:start_action_task()
    if self.action_task then
        return
    end
    self.action_task = true
    if not self.init_fork_vastai_info then
        skynet.fork(function ()
            self:change_vastai_info()
        end)
    end
end
 
return c_vastai_gpu