//
// Do NOT modify or remove this copyright and license
//
// Copyright (c) 2024 Seagate Technology LLC and/or its Affiliates, All Rights Reserved
//
// This software is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// ******************************************************************************************
//
// \file realtek_nvme_helper.h
// \brief Defines the functions for Realtek NVMe-USB pass-through

// All code in this file is from a Realtek USB to NVMe product specification for pass-through nvme commands.
// This code should only be used on products that are known to use this pass-through interface.

#include "realtek_nvme_helper.h"
#include "scsi_helper_func.h" //for ability to send a SCSI IO

eReturnValues build_Realtek_NVMe_CDB_And_Payload(uint8_t cdb[M_NONNULL_ARRAY REALTEK_NVME_CDB_SIZE],
                                                 eDataTransferDirection* M_NONNULL cdbDataDirection,
                                                 uint8_t* M_NULLABLE               dataPtr,
                                                 uint32_t                          dataSize,
                                                 eRealtekNVMCMDPhase               phase,
                                                 nvmeCmdCtx* M_NONNULL             nvmCmd)
{
    // static array should force passing a non-null pointer, but checking anyways since MSVC doesn't support marking
    // static arrays
    DISABLE_NONNULL_COMPARE
    if (cdb == M_NULLPTR)
    {
        return BAD_PARAMETER;
    }
    RESTORE_NONNULL_COMPARE

    safe_memset(cdb, REALTEK_NVME_CDB_SIZE, 0, REALTEK_NVME_CDB_SIZE);

    switch (phase)
    {
    case REALTEK_PHASE_COMMAND:
        *cdbDataDirection = XFER_DATA_OUT;
        if (dataPtr == M_NULLPTR || dataSize < REALTEK_NVME_CMD_PAYLOAD_LEN || nvmCmd == M_NULLPTR)
        {
            return BAD_PARAMETER;
        }
        else
        {
            safe_memset(dataPtr, REALTEK_NVME_CMD_PAYLOAD_LEN, 0, REALTEK_NVME_CMD_PAYLOAD_LEN);
            cdb[CDB_OPERATION_CODE] = REALTEK_NVME_PT_OPCODE_OUT;
            // set length to 0x40 - aka 64 bytes
            cdb[CDB_1] = M_Byte0(REALTEK_NVME_CMD_PAYLOAD_LEN); // length
            cdb[CDB_2] = M_Byte1(REALTEK_NVME_CMD_PAYLOAD_LEN); // length
            cdb[CDB_3] = C_CAST(uint8_t, REALTEK_PHASE_COMMAND);
            if (nvmCmd->commandType == NVM_ADMIN_CMD)
            {
                // set the cmd set field cdb 4-7
                cdb[CDB_4] = M_Byte0(REALTEK_NVME_ADMIN_CMD);
                cdb[CDB_5] = M_Byte1(REALTEK_NVME_ADMIN_CMD);
                cdb[CDB_6] = M_Byte2(REALTEK_NVME_ADMIN_CMD);
                cdb[CDB_7] = M_Byte3(REALTEK_NVME_ADMIN_CMD);
            }
            else
            {
                // set the cmd set field cdb 4-7
                cdb[CDB_4] = M_Byte0(REALTEK_NVME_IO_CMD);
                cdb[CDB_5] = M_Byte1(REALTEK_NVME_IO_CMD);
                cdb[CDB_6] = M_Byte2(REALTEK_NVME_IO_CMD);
                cdb[CDB_7] = M_Byte3(REALTEK_NVME_IO_CMD);
            }
            // Now setup the remaining command fields.
            // CDW0 is bytes 3:0
            dataPtr[0] = M_Byte0(nvmCmd->cmd.dwords.cdw0);
            dataPtr[1] = M_Byte1(nvmCmd->cmd.dwords.cdw0);
            dataPtr[2] = M_Byte2(nvmCmd->cmd.dwords.cdw0);
            dataPtr[3] = M_Byte3(nvmCmd->cmd.dwords.cdw0);
            // NSID is 7:4
            dataPtr[4] = M_Byte0(nvmCmd->cmd.dwords.cdw1);
            dataPtr[5] = M_Byte1(nvmCmd->cmd.dwords.cdw1);
            dataPtr[6] = M_Byte2(nvmCmd->cmd.dwords.cdw1);
            dataPtr[7] = M_Byte3(nvmCmd->cmd.dwords.cdw1);
            // CDW2 11:8
            dataPtr[8]  = M_Byte0(nvmCmd->cmd.dwords.cdw2);
            dataPtr[9]  = M_Byte1(nvmCmd->cmd.dwords.cdw2);
            dataPtr[10] = M_Byte2(nvmCmd->cmd.dwords.cdw2);
            dataPtr[11] = M_Byte3(nvmCmd->cmd.dwords.cdw2);
            // CDW3 15:12
            dataPtr[12] = M_Byte0(nvmCmd->cmd.dwords.cdw3);
            dataPtr[13] = M_Byte1(nvmCmd->cmd.dwords.cdw3);
            dataPtr[14] = M_Byte2(nvmCmd->cmd.dwords.cdw3);
            dataPtr[15] = M_Byte3(nvmCmd->cmd.dwords.cdw3);
            // metadata ptr is
            dataPtr[16] = M_Byte0(nvmCmd->cmd.dwords.cdw4);
            dataPtr[17] = M_Byte1(nvmCmd->cmd.dwords.cdw4);
            dataPtr[18] = M_Byte2(nvmCmd->cmd.dwords.cdw4);
            dataPtr[19] = M_Byte3(nvmCmd->cmd.dwords.cdw4);
            dataPtr[20] = M_Byte0(nvmCmd->cmd.dwords.cdw5);
            dataPtr[21] = M_Byte1(nvmCmd->cmd.dwords.cdw5);
            dataPtr[22] = M_Byte2(nvmCmd->cmd.dwords.cdw5);
            dataPtr[23] = M_Byte3(nvmCmd->cmd.dwords.cdw5);
            // data ptr is
            dataPtr[24] = M_Byte0(nvmCmd->cmd.dwords.cdw6);
            dataPtr[25] = M_Byte1(nvmCmd->cmd.dwords.cdw6);
            dataPtr[26] = M_Byte2(nvmCmd->cmd.dwords.cdw6);
            dataPtr[27] = M_Byte3(nvmCmd->cmd.dwords.cdw6);
            dataPtr[28] = M_Byte0(nvmCmd->cmd.dwords.cdw7);
            dataPtr[29] = M_Byte1(nvmCmd->cmd.dwords.cdw7);
            dataPtr[30] = M_Byte2(nvmCmd->cmd.dwords.cdw7);
            dataPtr[31] = M_Byte3(nvmCmd->cmd.dwords.cdw7);
            dataPtr[32] = M_Byte0(nvmCmd->cmd.dwords.cdw8);
            dataPtr[33] = M_Byte1(nvmCmd->cmd.dwords.cdw8);
            dataPtr[34] = M_Byte2(nvmCmd->cmd.dwords.cdw8);
            dataPtr[35] = M_Byte3(nvmCmd->cmd.dwords.cdw8);
            dataPtr[36] = M_Byte0(nvmCmd->cmd.dwords.cdw9);
            dataPtr[37] = M_Byte1(nvmCmd->cmd.dwords.cdw9);
            dataPtr[38] = M_Byte2(nvmCmd->cmd.dwords.cdw9);
            dataPtr[39] = M_Byte3(nvmCmd->cmd.dwords.cdw9);
            // CDW10 is 43:40
            dataPtr[40] = M_Byte0(nvmCmd->cmd.dwords.cdw10);
            dataPtr[41] = M_Byte1(nvmCmd->cmd.dwords.cdw10);
            dataPtr[42] = M_Byte2(nvmCmd->cmd.dwords.cdw10);
            dataPtr[43] = M_Byte3(nvmCmd->cmd.dwords.cdw10);
            // CDW11 is 47:44
            dataPtr[44] = M_Byte0(nvmCmd->cmd.dwords.cdw11);
            dataPtr[45] = M_Byte1(nvmCmd->cmd.dwords.cdw11);
            dataPtr[46] = M_Byte2(nvmCmd->cmd.dwords.cdw11);
            dataPtr[47] = M_Byte3(nvmCmd->cmd.dwords.cdw11);
            // CDW12 is 51:48
            dataPtr[48] = M_Byte0(nvmCmd->cmd.dwords.cdw12);
            dataPtr[49] = M_Byte1(nvmCmd->cmd.dwords.cdw12);
            dataPtr[50] = M_Byte2(nvmCmd->cmd.dwords.cdw12);
            dataPtr[51] = M_Byte3(nvmCmd->cmd.dwords.cdw12);
            // CDW13 is 55:52
            dataPtr[52] = M_Byte0(nvmCmd->cmd.dwords.cdw13);
            dataPtr[53] = M_Byte1(nvmCmd->cmd.dwords.cdw13);
            dataPtr[54] = M_Byte2(nvmCmd->cmd.dwords.cdw13);
            dataPtr[55] = M_Byte3(nvmCmd->cmd.dwords.cdw13);
            // CDW14 is 59:56
            dataPtr[56] = M_Byte0(nvmCmd->cmd.dwords.cdw14);
            dataPtr[57] = M_Byte1(nvmCmd->cmd.dwords.cdw14);
            dataPtr[58] = M_Byte2(nvmCmd->cmd.dwords.cdw14);
            dataPtr[59] = M_Byte3(nvmCmd->cmd.dwords.cdw14);
            // CDW15 is 63:60
            dataPtr[60] = M_Byte0(nvmCmd->cmd.dwords.cdw15);
            dataPtr[61] = M_Byte1(nvmCmd->cmd.dwords.cdw15);
            dataPtr[62] = M_Byte2(nvmCmd->cmd.dwords.cdw15);
            dataPtr[63] = M_Byte3(nvmCmd->cmd.dwords.cdw15);
        }
        break;
    case REALTEK_PHASE_DATA:
        if (nvmCmd == M_NULLPTR)
        {
            return BAD_PARAMETER;
        }
        // read uses in, write and non-data use out
        if (nvmCmd->commandDirection == XFER_DATA_IN)
        {
            cdb[CDB_OPERATION_CODE] = REALTEK_NVME_PT_OPCODE_IN;
            *cdbDataDirection       = XFER_DATA_IN;
            cdb[CDB_6]              = M_Byte0(REALTEK_DATA_IN);
            cdb[CDB_7]              = M_Byte1(REALTEK_DATA_IN);
        }
        else
        {
            cdb[CDB_OPERATION_CODE] = REALTEK_NVME_PT_OPCODE_OUT;
            *cdbDataDirection       = XFER_DATA_OUT;
            cdb[CDB_6]              = M_Byte0(REALTEK_DATA_OUT);
            cdb[CDB_7]              = M_Byte1(REALTEK_DATA_OUT);
        }
        // length of transfer can be zero to 256KiB
        cdb[CDB_1] = M_Byte0(nvmCmd->dataSize); // length
        cdb[CDB_2] = M_Byte1(nvmCmd->dataSize); // length
        cdb[CDB_4] = M_Byte2(nvmCmd->dataSize); // length
        cdb[CDB_5] = M_Byte3(nvmCmd->dataSize); // length
        if (nvmCmd->dataSize == UINT32_C(0))
        {
            *cdbDataDirection = XFER_NO_DATA;
            cdb[CDB_6]        = M_Byte0(REALTEK_NO_DATA);
            cdb[CDB_7]        = M_Byte1(REALTEK_NO_DATA);
        }
        cdb[CDB_3] = C_CAST(uint8_t, REALTEK_PHASE_DATA);
        break;
    case REALTEK_PHASE_COMPLETION:
        *cdbDataDirection       = XFER_DATA_IN;
        cdb[CDB_OPERATION_CODE] = REALTEK_NVME_PT_OPCODE_IN;
        // length must be 16 bytes
        cdb[CDB_1] = M_Byte0(REALTEK_NVME_COMPLETION_PAYLOAD_LEN);
        cdb[CDB_2] = M_Byte1(REALTEK_NVME_COMPLETION_PAYLOAD_LEN);
        cdb[CDB_3] = C_CAST(uint8_t, REALTEK_PHASE_COMPLETION);
        break;
    default:
        return BAD_PARAMETER;
    }

    return SUCCESS;
}

eReturnValues send_Realtek_NVMe_Cmd(nvmeCmdCtx* nvmCmd)
{
    eReturnValues ret = SUCCESS;
    DECLARE_ZERO_INIT_ARRAY(uint8_t, realtekCDB, REALTEK_NVME_CDB_SIZE);
    DECLARE_ZERO_INIT_ARRAY(uint8_t, realtekPayload, REALTEK_NVME_CMD_PAYLOAD_LEN);
    eDataTransferDirection realtekCDBDir = 0;

    if (nvmCmd == M_NULLPTR)
    {
        return BAD_PARAMETER;
    }

    if (nvmCmd->dataSize > REALTEK_NVME_MAX_TRANSFER_SIZE_BYTES || nvmCmd->commandDirection == XFER_DATA_IN_OUT ||
        nvmCmd->commandDirection == XFER_DATA_OUT_IN)
    {
        return OS_COMMAND_NOT_AVAILABLE;
    }
    // 1. build CDB & data for command to send
    // Send CDB to set the command values that will be used to issue a command.
    ret = build_Realtek_NVMe_CDB_And_Payload(realtekCDB, &realtekCDBDir, realtekPayload, REALTEK_NVME_CMD_PAYLOAD_LEN,
                                             REALTEK_PHASE_COMMAND, nvmCmd);
    if (SUCCESS != ret)
    {
        return ret;
    }
    ret = scsi_Send_Cdb(nvmCmd->device, realtekCDB, REALTEK_NVME_CDB_SIZE, realtekPayload, REALTEK_NVME_CMD_PAYLOAD_LEN,
                        realtekCDBDir, NULL, 0, DEFAULT_COMMAND_TIMEOUT);
    if (SUCCESS != ret)
    {
        return ret;
    }

    // 2. build CDB for data trasfer (including non-data)
    // send CDB for data transfer (this tiggers the actual command to execute)
    safe_memset(realtekCDB, REALTEK_NVME_CDB_SIZE, 0, REALTEK_NVME_CDB_SIZE);
    ret = build_Realtek_NVMe_CDB_And_Payload(realtekCDB, &realtekCDBDir, nvmCmd->ptrData, nvmCmd->dataSize,
                                             REALTEK_PHASE_DATA, nvmCmd);
    if (SUCCESS != ret)
    {
        return ret;
    }
    eReturnValues sendRet = scsi_Send_Cdb(nvmCmd->device, realtekCDB, REALTEK_NVME_CDB_SIZE, nvmCmd->ptrData,
                                          nvmCmd->dataSize, realtekCDBDir, NULL, 0, nvmCmd->timeout);
    // NOTE: do not fail the command or anything else YET.
    // Need to request the response information from the command.
    // TODO: There may be some sense data outputs where the return response info won't work or isn't necessary, but they
    // don't seem documented today. Most likely only for illegal requests.
    if (sendRet != OS_COMMAND_TIMEOUT)
    {
        // 3. build CDB for response info
        // send CDB for response info
        safe_memset(realtekCDB, REALTEK_NVME_CDB_SIZE, 0, REALTEK_NVME_CDB_SIZE);
        safe_memset(realtekPayload, REALTEK_NVME_CMD_PAYLOAD_LEN, 0, REALTEK_NVME_CMD_PAYLOAD_LEN);
        ret = build_Realtek_NVMe_CDB_And_Payload(realtekCDB, &realtekCDBDir, realtekPayload,
                                                 REALTEK_NVME_COMPLETION_PAYLOAD_LEN, REALTEK_PHASE_COMPLETION, nvmCmd);
        if (SUCCESS == scsi_Send_Cdb(nvmCmd->device, realtekCDB, REALTEK_NVME_CDB_SIZE, realtekPayload,
                                     REALTEK_NVME_COMPLETION_PAYLOAD_LEN, realtekCDBDir, NULL, 0, 15))
        {
            nvmCmd->commandCompletionData.dw0Valid = true;
            nvmCmd->commandCompletionData.dw1Valid = true;
            nvmCmd->commandCompletionData.dw2Valid = true;
            nvmCmd->commandCompletionData.dw3Valid = true;
            // This is the correct endianness as it has been validated by trying commands with invalid fields to make
            // sure it is correctly interpretted.
            nvmCmd->commandCompletionData.dw0 =
                M_BytesTo4ByteValue(realtekPayload[3], realtekPayload[2], realtekPayload[1], realtekPayload[0]);
            nvmCmd->commandCompletionData.dw1 =
                M_BytesTo4ByteValue(realtekPayload[7], realtekPayload[6], realtekPayload[5], realtekPayload[4]);
            nvmCmd->commandCompletionData.dw2 =
                M_BytesTo4ByteValue(realtekPayload[11], realtekPayload[10], realtekPayload[9], realtekPayload[8]);
            nvmCmd->commandCompletionData.dw3 =
                M_BytesTo4ByteValue(realtekPayload[15], realtekPayload[14], realtekPayload[13], realtekPayload[12]);
        }
    }
    return ret;
}

// cmds supported:
// admin:
//  - get log page (2)
//  - identify (6)
//  - security send/receive
//  - get features (no way to get completion registers at this time)
//  - format
// CMDs to test:
//  - firmware image download - cdw 11 instead??? this is offset, whereas xfer size is cdw10 as dwords
//  - firmware image commit - action and slot at most

// NOT SUPPORTED:
//  device self-test
//  sanitize
//  set features

eReturnValues build_Realtek_Basic_NVMe_CDB_And_Payload(uint8_t     cdb[M_NONNULL_ARRAY REALTEK_NVME_CDB_SIZE],
                                                       nvmeCmdCtx* nvmCmd)
{
    eReturnValues ret = SUCCESS;
    if (nvmCmd->commandType != NVM_ADMIN_CMD || nvmCmd->dataSize > REALTEK_BASIC_MAX_TRANSFER_SIZE_BYTES)
    {
        return OS_COMMAND_NOT_AVAILABLE;
    }
    if ((nvmCmd->commandDirection == XFER_DATA_OUT && nvmCmd->dataSize > 0) || nvmCmd->commandDirection == XFER_NO_DATA)
    {
        // use for data out and no-data. No data tested with non-data command
        cdb[0] = REALTEK_NVME_PT_OPCODE_OUT;
    }
    else
    {
        cdb[0] = REALTEK_NVME_PT_OPCODE_IN;
    }
    cdb[1] = M_Byte0(nvmCmd->dataSize);
    cdb[2] = M_Byte1(nvmCmd->dataSize);
    cdb[3] = nvmCmd->cmd.adminCmd.opcode;
    // general rule that lowest byte of dword 0 goes here.
    // This seems to work for get log page, get features, identify commands - TJE
    cdb[4] = M_Byte0(nvmCmd->cmd.adminCmd.cdw10);
    if (nvmCmd->cmd.adminCmd.opcode == NVME_ADMIN_CMD_FORMAT_NVM)
    {
        // Note sure if CDB[5] gets used or not. It seems accepted for specifying which LBA format, but that is
        //  part of cdb[4]...not sure if SES can be specified or not?
        // cdb[5] = M_Byte1(nvmCmd->cmd.adminCmd.cdw10) >> 1;
    }
    return ret;
}

eReturnValues send_Realtek_Basic_NVMe_Cmd(nvmeCmdCtx* nvmCmd)
{
    eReturnValues ret = SUCCESS;
    DECLARE_ZERO_INIT_ARRAY(uint8_t, realtekCDB, REALTEK_NVME_CDB_SIZE);

    if (nvmCmd == M_NULLPTR)
    {
        return BAD_PARAMETER;
    }

    ret = build_Realtek_Basic_NVMe_CDB_And_Payload(realtekCDB, nvmCmd);
    if (SUCCESS != ret)
    {
        return ret;
    }
    eReturnValues sendRet = scsi_Send_Cdb(nvmCmd->device, realtekCDB, REALTEK_NVME_CDB_SIZE, nvmCmd->ptrData,
                                          nvmCmd->dataSize, nvmCmd->commandDirection, NULL, 0, nvmCmd->timeout);
    // TODO: Translate sense to NVMe status?
    // NOTE: There does not appear to be a way to read completion DWords from the adapter from what I can tell - TJE
    // NOTE: For unsupported commands over this passthrough, the sense data returns invalid command operation code,
    // meaning
    //       the NVMe opcode, not the SCSI E4/E5 opcode. -TJE
    return sendRet;
}
