// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "firmware_spi_updater.hpp"

#include "google3/host_commands.h"

#include "message_util.hpp"
#include "payload_update_interface.hpp"

#include <stdplus/handle/managed.hpp>
#include <stdplus/print.hpp>
#include <stdplus/raw.hpp>
#include <xyz/openbmc_project/Control/Hoth/error.hpp>

#include <algorithm>
#include <cmath>
#include <format>
#include <optional>
#include <span>
#include <string>

namespace google
{
namespace hoth
{
namespace internal
{
namespace
{
using boost::endian::big_uint24_t;
using boost::endian::big_uint32_t;
using sdbusplus::error::xyz::openbmc_project::control::hoth::FirmwareFailure;

// Each SPI erase size: 4096 bytes
constexpr uint16_t singleEraseSize = 4096;
// Each SPI write size: 256 bytes
[[maybe_unused]] constexpr uint16_t singleWriteSize = 256;
// Each Read SPI read size set to 1000 bytes (mailbox is 1024 bytes)
constexpr uint16_t singleReadSize = 1000;

constexpr uint16_t writeEnableMosiLen = 1;
constexpr uint16_t writeEnableMisoLen = 0;
constexpr uint16_t eraseMisoLen = 0;
constexpr uint16_t writeMisoLen = 0;

// SPI opcodes
constexpr uint8_t opWriteEnable = 0x06;
constexpr uint8_t opErase = 0x20;
constexpr uint8_t opPageWrite = 0x02;
constexpr uint8_t opFastRead = 0x0B;
constexpr uint8_t opEnter4B = 0xB7;
constexpr uint8_t opExit4B = 0xE9;

/** @brief SPI WriteEnable Transaction
 */
struct SpiWriteEnableTransaction
{
    EcSpiOperationHeader header = {writeEnableMosiLen, writeEnableMisoLen};
    uint8_t opcode = opWriteEnable;
};

/** @brief SPI Erase Transaction
 */
struct SpiEraseTransaction
{
    explicit SpiEraseTransaction(uint8_t addressSize)
    {
        header.mosiLen = 1 + addressSize;
        header.misoLen = eraseMisoLen;
    }

    EcSpiOperationHeader header;
    uint8_t opcode = opErase;
};

/** @brief SPI Write Transaction
 */
struct SpiWriteTransaction
{
    SpiWriteTransaction(uint8_t addressSize, uint16_t dataSize)
    {
        header.mosiLen = addressSize + dataSize + 1;
        header.misoLen = writeMisoLen;
    }

    EcSpiOperationHeader header;
    uint8_t opcode = opPageWrite;
};

/** @brief SPI Read Transaction
 */
struct SpiReadTransaction
{
    SpiReadTransaction(uint8_t addressSize, uint16_t readSize)
    {
        // SPI fast read returns 5 or 6 extra 0x00 bytes at the begining of
        // response
        const uint16_t spiFastReadPadding = 2 + addressSize;
        header.mosiLen = 1 + addressSize;
        header.misoLen = readSize + spiFastReadPadding;
    }

    EcSpiOperationHeader header;
    uint8_t opcode = opFastRead;
};

/** @brief SPI Enter 4-Byte Addressing Mode Transaction
 */
struct SpiEnter4BTransaction
{
    SpiEnter4BTransaction() : header({1, 0})
    {}
    EcSpiOperationHeader header;
    uint8_t opcode = opEnter4B;
};

/** @brief SPI Exit 4-Byte Addressing Mode Transaction
 */
struct SpiExit4BTransaction
{
    SpiExit4BTransaction() : header({1, 0})
    {}
    EcSpiOperationHeader header;
    uint8_t opcode = opExit4B;
};

/** @brief Get ec_host_response.result from Hoth response
 */
uint16_t getHothResponseResult(const std::vector<uint8_t>& response)
{
    return stdplus::raw::refFrom<RspHeader>(response).result;
}

/** @brief Append object to byte array
 */
template <typename T>
void appendToByteArray(std::vector<uint8_t>& arr, const T& t)
{
    const uint8_t* p = reinterpret_cast<const uint8_t*>(&t);
    arr.insert(arr.end(), p, p + sizeof(t));
}
} // namespace

FirmwareSpiUpdater::SpiWritePreparer::SpiWritePreparer(
    FirmwareSpiUpdater* updater, ResetMode mode) :
    fwUpdater(updater),
    resetMode(mode)
{
    if (!fwUpdater)
    {
        stdplus::print(stderr, "{}: Null firmware updater. Skip preparing.\n",
                       __func__);
        return;
    }
    if (resetMode == ResetMode::never)
    {
        return;
    }
    const char* enclosingFuncName = __func__;
    const std::optional<bool> sps_passthrough_mode =
        [this, &enclosingFuncName] -> std::optional<bool> {
        try
        {
            return fwUpdater->getSpsPassthrough();
        }
        catch (const std::exception& e)
        {
            stdplus::print(stderr,
                           "{}: Unable to check SPS passthrough status: {}. "
                           "Trying prepare: {}\n",
                           enclosingFuncName, e.what(), resetMode == ResetMode::ignore);
            return std::nullopt;
        }
    }();

    const bool assumed_sps_through_mode = resetMode != ResetMode::ignore;
    if (!sps_passthrough_mode.value_or(assumed_sps_through_mode))
    {
        return;
    }
    if (resetMode == ResetMode::needed_active)
    {
        if (fwUpdater->setSpsPassthrough(false))
        {
            setSpsPassThoughDisabled = true;
            return;
        }

        stdplus::print(stderr, "{}: Failed to disable SPS passthrough.\n",
                       __func__);
    }
    fwUpdater->resetTarget(EC_TARGET_RESET_OPTION_SET);
    didResetTarget = true;
}

FirmwareSpiUpdater::SpiWritePreparer::SpiWritePreparer(
    SpiWritePreparer &&preparer) noexcept
    : fwUpdater(preparer.fwUpdater),
      setSpsPassThoughDisabled(preparer.setSpsPassThoughDisabled),
      didResetTarget(preparer.didResetTarget), resetMode(preparer.resetMode)
{
    preparer.fwUpdater = nullptr;
    preparer.setSpsPassThoughDisabled = false;
    preparer.didResetTarget = false;
}

FirmwareSpiUpdater::SpiWritePreparer::~SpiWritePreparer()
{
    if (setSpsPassThoughDisabled && fwUpdater &&
        !fwUpdater->setSpsPassthrough(true))
    {
        stdplus::print(stderr, "{}: Failed to reenable SPS passthrough.\n",
                       __func__);
    }
    if (didResetTarget)
    {
        releaseTarget(auto(fwUpdater));
    }
}

FirmwareSpiUpdater::SpiWritePreparer FirmwareSpiUpdater::prepareSpiWrite()
{
    SpiWritePreparer preparer(this, targetReset);
    return preparer;
}

FirmwareSpiUpdater::FirmwareSpiUpdater(HostCommand* hostCmd,
                                       uint8_t addressSize,
                                       ResetMode targetReset,
                                       bool ignoreAddressMode) :
    hostCmd(hostCmd),
    addressSize(addressSize), targetReset(targetReset),
    ignoreAddressMode(ignoreAddressMode)
{
    if (addressSize != 3 && addressSize != 4)
    {
        throw std::runtime_error(
            std::format("Unexpected address size: {}", addressSize));
    }
}

uint32_t FirmwareSpiUpdater::supportUpdateWithHothSpi(uint32_t firmwareSize)
{
    get_region_request request{GET_REGION_HOTH_UPDATE};
    std::vector<uint8_t> response = hostCmd->sendCommand(
        EC_CMD_BOARD_SPECIFIC_BASE | EC_PRV_CMD_HOTH_GET_REGION_FROM_IMG_DESC,
        0, &request, sizeof(request));

    auto rsp = std::span<const uint8_t>(response);
    const auto& hdr = stdplus::raw::extractRef<RspHeader>(rsp);
    if (hdr.result != EC_RES_SUCCESS)
    {
        stdplus::print(stderr, "Hoth GET_REGION_HOTH_UPDATE result: {}\n",
                       static_cast<uint8_t>(hdr.result));
        throw FirmwareFailure();
    }

    const auto& updateRegionRsp =
        stdplus::raw::refFrom<HothUpdateRegionResponse>(rsp);
    if (updateRegionRsp.flags & HOTH_UPDATE_REGION_FLAG_VALID)
    {
        uint32_t requiredSize =
            singleEraseSize *
            ((firmwareSize + singleEraseSize - 1) / singleEraseSize);
        if (updateRegionRsp.len < requiredSize)
        {
            stdplus::print(
                stderr,
                "Insufficient flash storage! Requires: {}, available {}\n",
                requiredSize, static_cast<uint32_t>(updateRegionRsp.len));
            throw FirmwareFailure();
        }
        return updateRegionRsp.offset;
    }
    throw FirmwareFailure();
}

void FirmwareSpiUpdater::resetTarget(uint8_t option)
{
    ec_request_reset_target request;
    request.target_id = RESET_TARGET_ID_RSTCTRL0;
    request.reset_option = option;
    std::vector<uint8_t> response = hostCmd->sendCommand(
        EC_CMD_BOARD_SPECIFIC_BASE | EC_PRV_CMD_HOTH_RESET_TARGET, 0, &request,
        sizeof(request));
    auto rsp = std::span<const uint8_t>(response);
    const auto& hdr = stdplus::raw::extractRef<RspHeader>(rsp);
    if (hdr.result != EC_RES_SUCCESS)
    {
        stdplus::print(stderr, "Hoth EC_PRV_CMD_HOTH_RESET_TARGET result: {}\n",
                       static_cast<uint8_t>(hdr.result));
        throw FirmwareFailure();
    }
}

void FirmwareSpiUpdater::releaseTarget(FirmwareSpiUpdater*&& updater) noexcept
{
    // Catch exception here so it won't throw a second exception after
    // FirmwareSpiUpdater::update() throws an exception.
    try
    {
        updater->resetTarget(EC_TARGET_RESET_OPTION_RELEASE);
    }
    catch (const std::exception& e)
    {
        stdplus::print(stderr, "Error releasing VLC: {}\n", e.what());
    }
}

bool FirmwareSpiUpdater::getSpsPassthrough()
{
    uint8_t request{};

    std::vector<uint8_t> response =
        hostCmd->sendCommand(EC_CMD_BOARD_SPECIFIC_BASE |
                                 EC_PRV_CMD_HOTH_GET_SPS_PASSTHROUGH_STATUS,
                             0, &request, sizeof(request));

    auto rsp = std::span<const uint8_t>(response);
    const auto& hdr = stdplus::raw::extractRef<RspHeader>(rsp);
    if (hdr.result != EC_RES_SUCCESS)
    {
        stdplus::print(stderr, "Hoth GET_SPS_PASSTHROUGH_STATUS result: {}\n",
                       static_cast<uint8_t>(hdr.result));
        throw FirmwareFailure();
    }

    const auto& spsRsp =
        stdplus::raw::refFrom<SpsPassthroughStatusResponse>(rsp);
    return spsRsp.enabled;
}

bool FirmwareSpiUpdater::setSpsPassthrough(const bool enable)
{
    uint8_t request{};

    std::vector<uint8_t> response = hostCmd->sendCommand(
        EC_CMD_BOARD_SPECIFIC_BASE |
            (enable ? EC_PRV_CMD_HOTH_SPS_PASSTHROUGH_ENABLE
                    : EC_PRV_CMD_HOTH_SPS_PASSTHROUGH_DISABLE),
        0, &request, sizeof(request));
    auto rsp = std::span<const uint8_t>(response);
    const auto& hdr = stdplus::raw::extractRef<RspHeader>(rsp);
    if (hdr.result != EC_RES_SUCCESS)
    {
        stdplus::print(stderr, "{}: Hoth SPS_PASSTHROGH_{} result: {}\n",
                       __func__, enable ? "ENABLE" : "DISBALE",
                       static_cast<uint8_t>(hdr.result));
        return false;
    }
    return true;
}

void FirmwareSpiUpdater::update(std::vector<uint8_t> firmwareData)
{
    uint32_t flashOffset = supportUpdateWithHothSpi(firmwareData.size());
    spiWrite(flashOffset, firmwareData);
}

void FirmwareSpiUpdater::spiWrite(uint32_t address, std::vector<uint8_t> data)
{
    const auto writePreparer = prepareSpiWrite();
    doUpdate(data, address);
}

void FirmwareSpiUpdater::doUpdate(const std::vector<uint8_t>& data,
                                  uint32_t address)
{
    // Enter or exit 4-byte address mode depending on the address of choice
    std::vector<uint8_t> request;
    if (addressSize == 3)
    {
        SpiExit4BTransaction exit4BTx;
        appendToByteArray(request, exit4BTx);
    }
    else
    {
        SpiEnter4BTransaction enter4BTx;
        appendToByteArray(request, enter4BTx);
    }

    std::vector<uint8_t> response = hostCmd->sendCommand(
        EC_CMD_BOARD_SPECIFIC_BASE | EC_PRV_CMD_HOTH_SPI_OPERATION, 0,
        request.data(), request.size());
    if (uint16_t result = getHothResponseResult(response);
        result != EC_RES_SUCCESS)
    {
        const std::string opType = (addressSize == 3) ? "exit" : "enter";
        const std::string errMsg = std::format(
            "Could not {} 4B addressing mode. Result: {}", opType, result);
        if (!ignoreAddressMode)
        {
            throw std::runtime_error(errMsg);
        }
        stdplus::print(stderr, "Could not {} 4B addressing mode. Result: {}\n",
                       opType, result);
        stdplus::print(stderr, "Ignoring address mode change failure\n");
    }

    erase(data.size(), address);
    writeData(data, address);
    verifyData(data, address);
}

void FirmwareSpiUpdater::erase(size_t firmwareSize, uint32_t address)
{
    std::vector<uint8_t> request;
    for (uint32_t endAddress = address + firmwareSize; address < endAddress;
         address += singleEraseSize)
    {
        request.clear();
        // Write enable transaction
        SpiWriteEnableTransaction writeEnableTx;
        appendToByteArray(request, writeEnableTx);
        // Erase transaction
        SpiEraseTransaction eraseTx(addressSize);
        appendToByteArray(request, eraseTx);
        appendAddress(request, address);

        std::vector<uint8_t> response = hostCmd->sendCommand(
            EC_CMD_BOARD_SPECIFIC_BASE | EC_PRV_CMD_HOTH_SPI_OPERATION, 0,
            request.data(), request.size());
        if (uint16_t result = getHothResponseResult(response);
            result != EC_RES_SUCCESS)
        {
            throw std::runtime_error(
                std::format("Hoth SPI Erase with result: {}", result));
        }
    }
}

void FirmwareSpiUpdater::writeData(const std::vector<uint8_t>& firmwareData,
                                   uint32_t address)
{
    std::vector<uint8_t> request;
    for (size_t i = 0, sz = firmwareData.size(); i < sz;)
    {
        request.clear();
        uint32_t beginAddress = address;
        uint32_t endAddress = beginAddress | 0xff;
        auto payloadSize = std::min(1 + endAddress - beginAddress,
                                    static_cast<uint32_t>(sz - i));

        // Write enable transaction
        SpiWriteEnableTransaction writeEnableTx;
        appendToByteArray(request, writeEnableTx);
        // Write transaction
        SpiWriteTransaction writeTx(addressSize, payloadSize);
        appendToByteArray(request, writeTx);
        appendAddress(request, beginAddress);
        // Write payload
        request.insert(request.end(), firmwareData.begin() + static_cast<int64_t>(i),
                       firmwareData.begin() + static_cast<int64_t>(i) + payloadSize);

        std::vector<uint8_t> response = hostCmd->sendCommand(
            EC_CMD_BOARD_SPECIFIC_BASE | EC_PRV_CMD_HOTH_SPI_OPERATION, 0,
            request.data(), request.size());
        if (uint16_t result = getHothResponseResult(response);
            result != EC_RES_SUCCESS)
        {
            throw std::runtime_error(
                std::format("Hoth SPI Write with result: {}", result));
        }

        i += payloadSize;
        address = endAddress + 1;
    }
}

void FirmwareSpiUpdater::verifyData(const std::vector<uint8_t>& firmwareData,
                                    uint32_t address)
{
    const uint16_t spiFastReadPadding = 2 + addressSize;
    std::vector<uint8_t> request;
    for (size_t i = 0, sz = firmwareData.size(); i < sz;)
    {
        request.clear();
        uint16_t readSize =
            std::min(static_cast<size_t>(singleReadSize), sz - i);
        // Read transaction
        SpiReadTransaction readTx(addressSize, readSize);
        appendToByteArray(request, readTx);
        appendAddress(request, address);

        std::vector<uint8_t> response = hostCmd->sendCommand(
            EC_CMD_BOARD_SPECIFIC_BASE | EC_PRV_CMD_HOTH_SPI_OPERATION, 0,
            request.data(), request.size());
        if (uint16_t result = getHothResponseResult(response);
            result != EC_RES_SUCCESS)
        {
            throw std::runtime_error(
                std::format("Hoth SPI Read error with result: {}", result));
        }

        // Compare firmware with data read from flash
        for (size_t j = i; j < i + readSize; ++j)
        {
            if (firmwareData[j] !=
                response[sizeof(RspHeader) + spiFastReadPadding + j - i])
            {
                throw std::runtime_error("SPI read got unexpected data");
            }
        }
        address += readSize;
        i += readSize;
    }
}

void FirmwareSpiUpdater::appendAddress(std::vector<uint8_t> &array,
                                       uint32_t address) const
{
    switch (addressSize) {
    case 3:
        appendToByteArray(array, big_uint24_t(address));
        break;
    case 4:
        appendToByteArray(array, big_uint32_t(address));
        break;
    default:
        break;
    }
}
} // namespace internal
} // namespace hoth
} // namespace google