-- Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved.
-- 
-- this file licensed under the 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 skynet = require 'skynet'
local bs = require 'mc.bitstring'
local log = require 'mc.logging'
local std_smbus = require 'protocol_open.protocol.std_smbus'
local cont = require 'protocol_open.common.cont'

local smbus_postbox = class(std_smbus)

local nv_gpu_reg_byte = bs.new('<<data0, data1, data2, data3>>')
local nv_gpu_reg_word = bs.new('<<data:32>>')

local REG_BIT_WIDTH <const> = 32

local cmd_status <const> = {
    CMD_STATUS_NULL = 0x00,              -- 无效的状态
    CMD_STATUS_ERR_REQUEST = 0x01,       -- SMBus错误
    CMD_STATUS_ERR_OPCODE = 0x02,        -- 无效的操作码
    CMD_STATUS_ERR_ARG1 = 0x03,          -- 无效的参数1
    CMD_STATUS_ERR_ARG2 = 0x04,          -- 无效的参数2
    CMD_STATUS_ERR_DATA = 0x05,          -- 无效数据
    CMD_STATUS_ERR_MISC = 0x06,          -- 未知错误
    CMD_STATUS_ERR_I2C_ACCESS = 0x07,    -- I2C内部错误
    CMD_STATUS_ERR_NOT_SUPPORTED = 0x08, -- 请求参数不被支持
    CMD_STATUS_ERR_NOT_AVAILABLE = 0x09, -- 请求当前不可用
    CMD_STATUS_ERR_BUSY = 0x0a,          -- 请求正在被处理
    CMD_STATUS_ERR_AGAIN = 0x0b,         -- 没有足够的资源处理该请求，请重试
    CMD_STATUS_ACCEPTED = 0x1c,          -- 异步请求已被接受
    CMD_STATUS_INACTIVE = 0x1d,          -- 未激活，不处理任何请求
    CMD_STATUS_READY = 0x1e,             -- 就绪
    CMD_STATUS_SUCCESS = 0x1f,           -- 请求执行成功
}

local cmd_status_str <const> = {
    [cmd_status.CMD_STATUS_NULL] = 'No response',
    [cmd_status.CMD_STATUS_ERR_REQUEST] = 'ERR_REQUEST',
    [cmd_status.CMD_STATUS_ERR_OPCODE] = 'ERR_OPCODE',
    [cmd_status.CMD_STATUS_ERR_ARG1] = 'ERR_ARG1',
    [cmd_status.CMD_STATUS_ERR_ARG2] = 'ERR_ARG2',
    [cmd_status.CMD_STATUS_ERR_DATA] = 'ERR_DATA',
    [cmd_status.CMD_STATUS_ERR_MISC] = 'ERR_MISC',
    [cmd_status.CMD_STATUS_ERR_I2C_ACCESS] = 'ERR_I2C_ACCESS',
    [cmd_status.CMD_STATUS_ERR_NOT_SUPPORTED] = 'ERR_NOT_SUPPORTED',
    [cmd_status.CMD_STATUS_ERR_NOT_AVAILABLE] = 'ERR_NOT_AVAILABLE',
    [cmd_status.CMD_STATUS_ERR_BUSY] = 'ERR_BUSY',
    [cmd_status.CMD_STATUS_ERR_AGAIN] = 'ERR_AGAIN',
    [cmd_status.CMD_STATUS_ACCEPTED] = 'ACCEPTED',
    [cmd_status.CMD_STATUS_INACTIVE] = 'INACTIVE',
    [cmd_status.CMD_STATUS_READY] = 'READY',
    [cmd_status.CMD_STATUS_SUCCESS] = 'SUCCESS',
}

local convert_type <const> = {
    KEEP_DIGITAL = 0,
    KEEP_ASCII = 1,
    ARRAY_TO_ASCII_BIGENDIAN = 2,
    ARRAY_TO_ASCII_LITTLEENDIAN = 3,
    KEEP_ARRAY = 4,
}

local request_params_template <const> = {
    name = true,
    opcode = true,
    arg1 = true,
    arg2 = true,
    capa_dword_num = true,
    capa_mask = true,
    data_out = true,
    ex_data_out = true,
    data_len = true,
    data_type = true,
    data_mask = true,
}

function smbus_postbox:validate_request_params(req)
    for key in pairs(req) do
        if not request_params_template[key] then
            return false
        end
    end
    return true
end

function smbus_postbox:_write_nv_gpu_reg(data, offset)
    local ok = self.super.send(self, offset, '\x04' .. data)
    if not ok then
        log:debug('write data 0x%08x to 0x%02x failed', data, offset)
    end
    return ok
end

function smbus_postbox:_read_nv_gpu_reg(offset)
    local ok, data = self.super.receive(self, offset, 5)
    if not ok or not data then
        log:debug('read data from 0x%02x failed, %s', offset, data)
        return
    end

    if data:sub(1, 1):byte() ~= 4 then
        log:debug('read length=%s from 0x%02x is invalid', data[1], offset)
        return
    end

    return data:sub(2, 5)
end

function smbus_postbox:_intf(cmd, data_in, data_out, ex_data_out)
    local NV_CMD_REG <const> = 0x5c
    local NV_GPU_STATUS_MASK <const> = 0x1f
    local ok
    if data_in then
        ok = self:_write_nv_gpu_reg(data_in, self.data_reg_addr)
        if not ok then
            return nil, nil
        end
    end

    ok = self:_write_nv_gpu_reg(cmd, NV_CMD_REG)
    if not ok then
        return nil, nil
    end

    skynet.sleep(10)  -- 100毫秒

    local status = cmd_status.CMD_STATUS_NULL
    local count = 0
    local res
    repeat
        res = self:_read_nv_gpu_reg(NV_CMD_REG)
        if not res then
            return nil, nil
        end
        status = nv_gpu_reg_byte:unpack(res).data3 & NV_GPU_STATUS_MASK
        if status ~= cmd_status.CMD_STATUS_NULL then
            break
        end

        count = count + 1
        skynet.sleep(100)  -- 1000毫秒
    until status ~= cmd_status.CMD_STATUS_NULL or count > 20

    log:debug('cmd status = %s', cmd_status_str[status] or 'Unknown')

    res = nil
    if data_out and status == cmd_status.CMD_STATUS_SUCCESS then
        res = self:_read_nv_gpu_reg(self.data_reg_addr)
        if not res then
            return nil, nil
        end
    end

    local res_ex
    if ex_data_out and status == cmd_status.CMD_STATUS_SUCCESS then
        res_ex = self:_read_nv_gpu_reg(self.data_reg_addr + 1)
        if not res_ex then
            return res, nil
        end
    end

    return res, res_ex
end

local function assemble_smbpbi_cmd(opcode, arg1, arg2)
    local CMD_EXE_FLAG <const> = 0x80
    return nv_gpu_reg_byte:pack({data0 = opcode, data1 = arg1, data2 = arg2, data3 = CMD_EXE_FLAG})
end

local function get_digital_actual_value(raw_data, mask)
    local offset = 0
    local mask_temp = mask
    while mask_temp ~= 0 do
        if mask_temp & 1 == 1 then
            break
        end
        mask_temp = mask_temp >> 1
        offset = offset + 1
    end

    return (raw_data & mask) >> offset
end

function smbus_postbox:_gpu_action_digital_process(request)
    local cmd = assemble_smbpbi_cmd(request.opcode, request.arg1, request.arg2)
    local res, _ = self:_intf(cmd, nil, request.data_out, nil)
    if res then
        local value = get_digital_actual_value(nv_gpu_reg_word:unpack(res).data, request.data_mask)
        return tostring(value)
    end
end

local function convert_ascii_to_string(data, data_type)
    if data_type == convert_type.ARRAY_TO_ASCII_BIGENDIAN then
        local tmp = nv_gpu_reg_byte:unpack(data)
        if tmp then
            return string.format("%02x%02x%02x%02x", tmp.data0, tmp.data1, tmp.data2, tmp.data3)
        end
    elseif data_type == convert_type.ARRAY_TO_ASCII_LITTLEENDIAN then
        local tmp = nv_gpu_reg_byte:unpack(data)
        if tmp then
            return string.format("%02x%02x%02x%02x", tmp.data3, tmp.data2, tmp.data1, tmp.data0)
        end
    else
        return data
    end
end

function smbus_postbox:_gpu_action_ascii_process(request)
    if not request.data_len then
        return
    end

    local NV_GPU_PROPERTY_MAX_LEN <const> = 32
    local data_len = (request.data_len > NV_GPU_PROPERTY_MAX_LEN and NV_GPU_PROPERTY_MAX_LEN or request.data_len)
    local fetch_time = data_len // 4
    if data_len % 4 ~= 0 then
        fetch_time = fetch_time + 1
    end
    local cmd
    local res
    local action_data = {}
    for i = 1, fetch_time do
        cmd = assemble_smbpbi_cmd(request.opcode, request.arg1, i - 1)
        res = self:_intf(cmd, nil, request.data_out, nil)
        if res then
            action_data[i] = convert_ascii_to_string(res, request.data_type)
        else
            return
        end
    end

    local result = table.concat(action_data)
    if not self.cont_table[request.name] then
        self.cont_table[request.name] = cont.new(nil, result)
    end

    return self.cont_table[request.name]:get_debounced_val(result)
end

local function convert_nvlink_data_to_array(res, info_array_offset, item_bits, info_array, nvlink_num)
    local data = nv_gpu_reg_word:unpack(res).data
    local data_mask = 0xffffffff >> (REG_BIT_WIDTH - item_bits)
    local reg_offset = 0
    while reg_offset < REG_BIT_WIDTH do
        if info_array_offset > nvlink_num then
            break
        end
        table.insert(info_array, info_array_offset, (data >> reg_offset) & data_mask)
        info_array_offset = info_array_offset + 1
        reg_offset = reg_offset + item_bits
    end

    return info_array_offset
end

function smbus_postbox:_gpu_action_array_process(request)
    -- 此处ex_data_out作为nvlink_num传入
    local MAX_NVLINK_NUM <const> = 256
    if not request.ex_data_out or request.ex_data_out == 0 or request.ex_data_out > MAX_NVLINK_NUM or 
        not request.data_len then
        return
    end

    local read_bits = request.ex_data_out * request.data_len
    local BITS_PER_READ <const> = REG_BIT_WIDTH * 2
    local fetch_time = read_bits // BITS_PER_READ
    if read_bits % BITS_PER_READ ~= 0 then
        fetch_time = fetch_time + 1
    end

    local cmd
    local res
    local res_ex
    local info_array = {}
    local info_array_offset = 1
    for i = 1, fetch_time do
        cmd = assemble_smbpbi_cmd(request.opcode, request.arg1, i - 1)
        res, res_ex = self:_intf(cmd, nil, request.data_out, request.ex_data_out)
        if res then
            info_array_offset = convert_nvlink_data_to_array(res, info_array_offset, request.data_len, info_array, 
                request.ex_data_out)
        end
        if res_ex then
            info_array_offset = convert_nvlink_data_to_array(res_ex, info_array_offset, request.data_len, info_array, 
                request.ex_data_out)
        end
    end

    return info_array
end

function smbus_postbox:_send_request_internal(request)
    local NV_CAPABILITY_NUM <const> = 4
    if request.capa_dword_num > NV_CAPABILITY_NUM then
        return
    end

    local OPCODE_GET_CAP <const> = 0x01
    local cmd = assemble_smbpbi_cmd(OPCODE_GET_CAP, request.capa_dword_num, 0)
    local res, _ = self:_intf(cmd, nil, request.data_out, nil)
    if not res then
        return
    end

    local capability = nv_gpu_reg_word:unpack(res).data
    if capability & request.capa_mask ~= request.capa_mask then
        log:info('%s not supported. Capability%s=0x%08x', request.name, request.capa_dword_num, capability)
        return
    end

    return self.process_table[request.data_type](self, request)
end

function smbus_postbox:ctor(params)
    self.name = 'std_postbox'
    self.cont_table = {}
    self.data_reg_addr = params.data_reg_addr
    self.process_table = {
        [convert_type.KEEP_DIGITAL] = smbus_postbox._gpu_action_digital_process,
        [convert_type.KEEP_ASCII] = smbus_postbox._gpu_action_ascii_process,
        [convert_type.ARRAY_TO_ASCII_BIGENDIAN] = smbus_postbox._gpu_action_ascii_process,
        [convert_type.ARRAY_TO_ASCII_LITTLEENDIAN] = smbus_postbox._gpu_action_ascii_process,
        [convert_type.KEEP_ARRAY] = smbus_postbox._gpu_action_array_process,
    }
end

return smbus_postbox