#include "NVMeFeatureStore.hpp"

#include "nvme_features.pb.h"

#include <endian.h>

#include <boost/asio/async_result.hpp>
#include <boost/asio/post.hpp>
#include <phosphor-logging/lg2.hpp>

#include <algorithm>
#include <array>
#include <cstring>
#include <iomanip>
#include <map>
#include <sstream>
#include <tuple>
#include <utility>
#include <vector>

static constexpr uint32_t nvmeGetFeaturesCapSaveable = 0x1;
static constexpr uint32_t nvmeGetFeaturesCapChangeable = 0x4;

struct FeatureMeta
{
    uint8_t fid;            ///< NVMe Feature Identifier
    std::string protoClass; ///< Fully qualified Protobuf class name
    FeatureScope scope;     ///< Applicability scope
};

/**
 * @brief Map of P0 features defined in the design document.
 */
static const std::map<std::string, FeatureMeta> supportedFeatures = {
    {"Arbitration",
     {0x01, "google.gbmc.nvme.base.Arbitration", FeatureScope::Controller}},
    {"PowerManagement",
     {0x02, "google.gbmc.nvme.base.PowerManagement", FeatureScope::Controller}},
    {"TemperatureThreshold",
     {0x04, "google.gbmc.nvme.base.TemperatureThreshold",
      FeatureScope::Controller}},
    {"ErrorRecovery",
     {0x05, "google.gbmc.nvme.base.ErrorRecovery", FeatureScope::Namespace}},
    {"VolatileWriteCache",
     {0x06, "google.gbmc.nvme.base.VolatileWriteCache",
      FeatureScope::Controller}},
    {"NumberOfQueues",
     {0x07, "google.gbmc.nvme.base.NumberOfQueues", FeatureScope::Controller}},
    {"InterruptCoalescing",
     {0x08, "google.gbmc.nvme.base.InterruptCoalescing",
      FeatureScope::Controller}},
    {"InterruptVectorConfiguration",
     {0x09, "google.gbmc.nvme.base.InterruptVectorConfiguration",
      FeatureScope::Controller}},
    {"WriteAtomicityNormal",
     {0x0A, "google.gbmc.nvme.base.WriteAtomicityNormal",
      FeatureScope::Controller}},
    {"AsynchronousEventConfiguration",
     {0x0B, "google.gbmc.nvme.base.AsynchronousEventConfiguration",
      FeatureScope::Controller}},
    {"OcpLatencyMonitor",
     {0xC5, "google.gbmc.nvme.base.OcpLatencyMonitor",
      FeatureScope::Controller}},
};

/**
 * @brief OCP Latency Monitor Feature Buffer (FID 0xC5) as defined by OCP NVMe
 * Spec. Matches platforms/storage/libnvme/nvme_types.h.
 */
struct __attribute__((packed)) OcpLatencyMonitorFeatureBuffer
{
    uint16_t active_bucket_timer_threshold;
    uint8_t active_threshold_a;
    uint8_t active_threshold_b;
    uint8_t active_threshold_c;
    uint8_t active_threshold_d;
    uint16_t active_latency_config;
    uint8_t active_latency_minimum_window;
    uint16_t debug_log_trigger_enable;
    uint8_t discard_debug_log;
    uint8_t latency_monitor_feature_enable;
    std::array<uint8_t, 4083> reserved;
};
static_assert(sizeof(OcpLatencyMonitorFeatureBuffer) == 4096,
              "Size of OcpLatencyMonitorFeatureBuffer must be 4096");

template <typename T>
static T parseProto(const std::vector<uint8_t>& data)
{
    T msg;
    if (!msg.ParseFromArray(data.data(), static_cast<int>(data.size())))
    {
        throw sdbusplus::xyz::openbmc_project::Common::Error::InvalidArgument();
    }
    return msg;
}

// Helper for async_initiate
static std::tuple<nvme_ex_ptr, uint32_t, std::vector<uint8_t>>
    asyncAdminGetFeatures(boost::asio::io_context& io,
                          const std::shared_ptr<NVMeMiIntf>& nvmeIntf,
                          nvme_mi_ctrl_t ctrl,
                          const NVMeMiIntf::GetFeaturesRequest& req,
                          boost::asio::yield_context yield)
{
    return boost::asio::async_initiate<boost::asio::yield_context,
                                       void(nvme_ex_ptr, uint32_t,
                                            std::vector<uint8_t>)>(
        [&io](auto&& callback, const std::shared_ptr<NVMeMiIntf>& nvmeIntf,
              nvme_mi_ctrl_t ctrl, const NVMeMiIntf::GetFeaturesRequest& req) {
        auto sharedCb = std::make_shared<std::decay_t<decltype(callback)>>(
            std::forward<decltype(callback)>(callback));
        nvmeIntf->adminGetFeatures(
            ctrl, req,
            [sharedCb, &io](nvme_ex_ptr ex, uint32_t resp,
                            std::span<uint8_t> data) mutable {
            std::vector<uint8_t> dataCopy(data.begin(), data.end());
            boost::asio::post(io, [sharedCb = std::move(sharedCb), ex, resp,
                                   dataCopy = std::move(dataCopy)]() mutable {
                (*sharedCb)(ex, resp, dataCopy);
            });
        });
    },
        yield, nvmeIntf, ctrl, req);
}

static std::tuple<nvme_ex_ptr, uint32_t> asyncAdminSetFeatures(
    boost::asio::io_context& io, const std::shared_ptr<NVMeMiIntf>& nvmeIntf,
    nvme_mi_ctrl_t ctrl, const NVMeMiIntf::SetFeaturesRequest& req,
    boost::asio::yield_context yield)
{
    return boost::asio::async_initiate<boost::asio::yield_context,
                                       void(nvme_ex_ptr, uint32_t)>(
        [&io](auto&& callback, const std::shared_ptr<NVMeMiIntf>& nvmeIntf,
              nvme_mi_ctrl_t ctrl, const NVMeMiIntf::SetFeaturesRequest& req) {
        auto sharedCb = std::make_shared<std::decay_t<decltype(callback)>>(
            std::forward<decltype(callback)>(callback));
        nvmeIntf->adminSetFeatures(
            ctrl, req, [sharedCb, &io](nvme_ex_ptr ex, uint32_t resp) mutable {
            boost::asio::post(io, [sharedCb = std::move(sharedCb), ex,
                                   resp]() mutable { (*sharedCb)(ex, resp); });
        });
    },
        yield, nvmeIntf, ctrl, req);
}

static std::tuple<nvme_ex_ptr, std::vector<uint8_t>> asyncAdminIdentify(
    boost::asio::io_context& io, const std::shared_ptr<NVMeMiIntf>& nvmeIntf,
    nvme_mi_ctrl_t ctrl, nvme_identify_cns cns, uint32_t nsid, uint16_t cntid,
    boost::asio::yield_context yield)
{
    return boost::asio::async_initiate<boost::asio::yield_context,
                                       void(nvme_ex_ptr, std::vector<uint8_t>)>(
        [&io](auto&& callback, const std::shared_ptr<NVMeMiIntf>& nvmeIntf,
              nvme_mi_ctrl_t ctrl, nvme_identify_cns cns, uint32_t nsid,
              uint16_t cntid) {
        auto sharedCb = std::make_shared<std::decay_t<decltype(callback)>>(
            std::forward<decltype(callback)>(callback));
        nvmeIntf->adminIdentify(
            ctrl, cns, nsid, cntid,
            [sharedCb, &io](nvme_ex_ptr ex, std::span<uint8_t> data) mutable {
            std::vector<uint8_t> dataCopy(data.begin(), data.end());
            boost::asio::post(io, [sharedCb = std::move(sharedCb), ex,
                                   dataCopy = std::move(dataCopy)]() mutable {
                (*sharedCb)(ex, dataCopy);
            });
        });
    },
        yield, nvmeIntf, ctrl, cns, nsid, cntid);
}

static std::tuple<std::error_code, std::vector<uint8_t>> asyncAdminGetLogPage(
    boost::asio::io_context& io, const std::shared_ptr<NVMeMiIntf>& nvmeIntf,
    nvme_mi_ctrl_t ctrl, nvme_cmd_get_log_lid lid, uint32_t nsid, uint8_t lsp,
    uint16_t lsi, boost::asio::yield_context yield)
{
    return boost::asio::async_initiate<boost::asio::yield_context,
                                       void(std::error_code,
                                            std::vector<uint8_t>)>(
        [&io](auto&& callback, const std::shared_ptr<NVMeMiIntf>& nvmeIntf,
              nvme_mi_ctrl_t ctrl, nvme_cmd_get_log_lid lid, uint32_t nsid,
              uint8_t lsp, uint16_t lsi) {
        auto sharedCb = std::make_shared<std::decay_t<decltype(callback)>>(
            std::forward<decltype(callback)>(callback));
        nvmeIntf->adminGetLogPage(
            ctrl, lid, nsid, lsp, lsi,
            [sharedCb, &io](const std::error_code& ec,
                            std::span<uint8_t> data) mutable {
            std::vector<uint8_t> dataCopy(data.begin(), data.end());
            boost::asio::post(io, [sharedCb = std::move(sharedCb), ec,
                                   dataCopy = std::move(dataCopy)]() mutable {
                (*sharedCb)(ec, dataCopy);
            });
        });
    },
        yield, nvmeIntf, ctrl, lid, nsid, lsp, lsi);
}

FeatureStore::FeatureStore(
    boost::asio::io_context& io, sdbusplus::asio::object_server& objServer,
    const std::shared_ptr<sdbusplus::asio::connection>& conn,
    const std::string& objectPath, std::shared_ptr<NVMeMiIntf> nvmeIntf,
    nvme_mi_ctrl_t ctrl, uint32_t nsid, FeatureScope scope) :
    io(io), objServer(objServer), conn(conn), objectPath(objectPath),
    nvmeIntf(std::move(nvmeIntf)), ctrl(ctrl), nsid(nsid), scope(scope)
{}

void FeatureStore::registerDbusInterface()
{
    dbusIntf = objServer.add_interface(objectPath,
                                       "xyz.openbmc_project.NVMe.FeatureStore");

    dbusIntf->register_property("ProtoLibraryVersion", std::string("2.0"),
                                sdbusplus::asio::PropertyPermission::readOnly);
    dbusIntf->register_property("FeatureCollection", featureCollection,
                                sdbusplus::asio::PropertyPermission::readOnly);
    dbusIntf->register_method(
        "GetFeature",
        [weakSelf = weak_from_this()](const boost::asio::yield_context& yield,
                                      const std::string& featureName) {
        auto self = weakSelf.lock();
        if (!self)
        {
            throw sdbusplus::xyz::openbmc_project::Common::Error::Unavailable();
        }
        return self->getFeature(yield, featureName);
    });

    dbusIntf->register_method(
        "SetFeature",
        [weakSelf = weak_from_this()](const boost::asio::yield_context& yield,
                                      const std::string& featureName,
                                      const std::vector<uint8_t>& data) {
        auto self = weakSelf.lock();
        if (!self)
        {
            throw sdbusplus::xyz::openbmc_project::Common::Error::Unavailable();
        }
        self->setFeature(yield, featureName, data);
    });

    dbusIntf->initialize();
}

void FeatureStore::init(std::function<void(const std::error_code&)>&& cb)
{
    (void)boost::asio::spawn(
        io, [self{shared_from_this()}, cb{std::move(cb)}](
                const boost::asio::yield_context& yield) mutable {
        if (!self->fetchIdentifyData(yield))
        {
            cb(std::make_error_code(std::errc::no_such_device));
            return;
        }
        self->discoverFeatures(yield);
        if (self->featureCollection.empty())
        {
            lg2::info(
                "[{PATH}]No NVMe features supported, skipping FeatureStore registration",
                "PATH", self->objectPath);
            cb(std::make_error_code(std::errc::not_supported));
            return;
        }
        self->registerDbusInterface();
        cb({});
    });
}

bool FeatureStore::fetchIdentifyData(const boost::asio::yield_context& yield)
{
    std::string path = this->objectPath;
    auto [ex, data] = asyncAdminIdentify(
        io, nvmeIntf, ctrl, NVME_IDENTIFY_CNS_CTRL, NVME_NSID_NONE, 0, yield);

    if (ex)
    {
        lg2::error(
            "[{PATH}]Identify Controller failed during FeatureStore init: {ERROR}",
            "PATH", path, "ERROR", ex->what());
        return false;
    }

    if (data.size() < sizeof(nvme_id_ctrl))
    {
        lg2::error("[{PATH}]Identify Controller returned insufficient data",
                   "PATH", path);
        return false;
    }
    nvme_id_ctrl idCtrl = {};
    std::memcpy(&idCtrl, data.data(), sizeof(nvme_id_ctrl));
    vwcSupported = (idCtrl.vwc & NVME_CTRL_VWC_PRESENT) != 0;
    oncs = le16toh(idCtrl.oncs);

    if (!vwcSupported)
    {
        lg2::info(
            "[{PATH}]Volatile Write Cache (FID 0x06) not supported by hardware",
            "PATH", path);
    }

    return true;
}

void FeatureStore::discoverFeatures(const boost::asio::yield_context& yield)
{
    std::string path = this->objectPath;
    bool oncsBit4 = (oncs & (1 << 4)) != 0;

    featureCollection.clear();

    std::vector<std::pair<std::string, FeatureMeta>> featuresToCheck;
    for (const auto& [name, meta] : supportedFeatures)
    {
        if (meta.scope == scope)
        {
            if (name == "VolatileWriteCache" && !vwcSupported)
            {
                continue;
            }
            featuresToCheck.emplace_back(name, meta);
        }
    }

    // Fetch Log Page 12h (Feature Identifiers Supported and Effects)
    std::vector<uint8_t> fidSupportedEffects;
    auto [ec, logData] = asyncAdminGetLogPage(
        io, nvmeIntf, ctrl, NVME_LOG_LID_FID_SUPPORTED_EFFECTS, NVME_NSID_NONE,
        0, 0, yield);
    if (!ec)
    {
        if (logData.size() == sizeof(struct nvme_fid_supported_effects_log))
        {
            fidSupportedEffects = std::move(logData);
        }
        else
        {
            lg2::error(
                "[{PATH}]Log Page 12h size mismatch: expected {EXPECTED}, got {GOT}",
                "PATH", path, "EXPECTED",
                sizeof(struct nvme_fid_supported_effects_log), "GOT",
                logData.size());
        }
    }
    else
    {
        lg2::error("[{PATH}]Failed to fetch Log Page 12h: {ERROR}", "PATH",
                   path, "ERROR", ec.message());
    }

    for (const auto& [name, meta] : featuresToCheck)
    {
        if (!fidSupportedEffects.empty())
        {
            size_t offset = static_cast<size_t>(meta.fid) * 4;
            if (offset + sizeof(uint32_t) <= fidSupportedEffects.size())
            {
                uint32_t entry = 0;
                std::memcpy(&entry, fidSupportedEffects.data() + offset,
                            sizeof(entry));
                entry = le32toh(entry);
                if ((entry & 0x1) == 0) // FSUPP bit
                {
                    lg2::info(
                        "[{PATH}]Feature {NAME} (FID {FID}) not supported per Log Page 12h",
                        "PATH", path, "NAME", name, "FID", meta.fid);
                    continue;
                }
            }
            else
            {
                lg2::error("[{PATH}]Log Page 12h data too short for FID {FID}",
                           "PATH", path, "FID", meta.fid);
                continue;
            }
        }
        std::stringstream ss;
        ss << "0x" << std::setfill('0') << std::setw(2) << std::hex
           << static_cast<int>(meta.fid);
        std::string fidStr = ss.str();

        if (!oncsBit4)
        {
            NVMeMiIntf::GetFeaturesRequest args{};
            args.fid = meta.fid;
            args.sel = NVME_GET_FEATURES_SEL_CURRENT; // 000b
            args.nsid = (scope == FeatureScope::Namespace) ? nsid
                                                           : NVME_NSID_NONE;
            args.cdw11 = 0;
            args.data_len = 0;

            auto [ex, cdw0, data] = asyncAdminGetFeatures(io, nvmeIntf, ctrl,
                                                          args, yield);

            if (ex)
            {
                lg2::info(
                    "[{PATH}]Feature {NAME} (FID {FID}) get feature query failed (unsupported): {ERROR}",
                    "PATH", path, "NAME", name, "FID", meta.fid, "ERROR",
                    ex->what());
                continue;
            }

            featureCollection.emplace_back(name + "_Current", fidStr,
                                           meta.protoClass);
            continue;
        }

        NVMeMiIntf::GetFeaturesRequest args{};
        args.fid = meta.fid;
        args.sel = NVME_GET_FEATURES_SEL_SUPPORTED; // 011b
        args.nsid = (scope == FeatureScope::Namespace) ? nsid : NVME_NSID_NONE;
        args.cdw11 = 0;
        args.data_len = 0;

        auto [ex, cdw0, data] = asyncAdminGetFeatures(io, nvmeIntf, ctrl, args,
                                                      yield);

        if (ex)
        {
            lg2::info(
                "[{PATH}]Feature {NAME} (FID {FID}) capabilities query failed (unsupported): {ERROR}",
                "PATH", path, "NAME", name, "FID", meta.fid, "ERROR",
                ex->what());
            continue;
        }

        featureCaps[meta.fid] = cdw0;

        featureCollection.emplace_back(name + "_Current", fidStr,
                                       meta.protoClass);

        if ((cdw0 & nvmeGetFeaturesCapChangeable) != 0)
        {
            featureCollection.emplace_back(name + "_Default", fidStr,
                                           meta.protoClass);
        }

        if ((cdw0 & nvmeGetFeaturesCapSaveable) != 0)
        {
            featureCollection.emplace_back(name + "_Saved", fidStr,
                                           meta.protoClass);
        }
    }
}

bool FeatureStore::isFeatureSupported(const std::string& featureName) const
{
    return std::any_of(featureCollection.begin(), featureCollection.end(),
                       [&featureName](const auto& item) {
        return std::get<0>(item) == featureName;
    });
}

size_t FeatureStore::getFeatureDataLen(uint8_t fid)
{
    switch (fid)
    {
        case 0x01: // Arbitration
        case 0x02: // PowerManagement
        case 0x04: // TemperatureThreshold
        case 0x05: // ErrorRecovery
        case 0x06: // VolatileWriteCache
        case 0x07: // NumberOfQueues
        case 0x08: // InterruptCoalescing
        case 0x09: // InterruptVectorConfiguration
        case 0x0A: // WriteAtomicityNormal
        case 0x0B: // AsynchronousEventConfiguration
            return 0;
        case 0xC5: // OcpLatencyMonitor
            return 4096;
        default:
            lg2::error("Unsupported feature ID for data length: {FID}", "FID",
                       lg2::hex, fid);
            throw sdbusplus::xyz::openbmc_project::Common::Error::
                InvalidArgument();
    }
}

std::vector<uint8_t>
    FeatureStore::getFeature(const boost::asio::yield_context& yield,
                             const std::string& featureName)
{
    std::string path = this->objectPath;
    if (!isFeatureSupported(featureName))
    {
        throw sdbusplus::xyz::openbmc_project::Common::Error::InvalidArgument();
    }
    size_t lastUnderscore = featureName.find_last_of('_');
    if (lastUnderscore == std::string::npos)
    {
        throw sdbusplus::xyz::openbmc_project::Common::Error::InvalidArgument();
    }

    std::string baseName = featureName.substr(0, lastUnderscore);
    std::string selectStr = featureName.substr(lastUnderscore + 1);

    auto it = supportedFeatures.find(baseName);
    if (it == supportedFeatures.end() || it->second.scope != scope)
    {
        throw sdbusplus::xyz::openbmc_project::Common::Error::InvalidArgument();
    }

    uint8_t fid = it->second.fid;
    uint8_t sel = 0;
    if (selectStr == "Current")
    {
        sel = 0;
    }
    else if (selectStr == "Default")
    {
        sel = 1;
    }
    else if (selectStr == "Saved")
    {
        sel = 2;
    }
    else
    {
        throw sdbusplus::xyz::openbmc_project::Common::Error::InvalidArgument();
    }

    NVMeMiIntf::GetFeaturesRequest args{};
    args.fid = fid;
    args.sel = sel;
    args.nsid = (scope == FeatureScope::Namespace) ? nsid : NVME_NSID_NONE;
    args.cdw11 = 0;
    args.data_len = getFeatureDataLen(fid);

    auto [ex, cdw0, data] = asyncAdminGetFeatures(io, nvmeIntf, ctrl, args,
                                                  yield);

    if (ex)
    {
        lg2::error("[{PATH}]NVMe Get Features (FID {FID}) failed: {ERROR}",
                   "PATH", path, "FID", fid, "ERROR", ex->what());
        throw sdbusplus::xyz::openbmc_project::Common::Error::Unavailable();
    }

    std::vector<uint8_t> protoData;

    switch (fid)
    {
        case 0x01:
        {
            google::gbmc::nvme::base::Arbitration msg;
            msg.set_arbitration_burst(cdw0 & 0xFF);
            msg.set_low_priority_weight((cdw0 >> 8) & 0xFF);
            msg.set_medium_priority_weight((cdw0 >> 16) & 0xFF);
            msg.set_high_priority_weight((cdw0 >> 24) & 0xFF);
            protoData.resize(msg.ByteSizeLong());
            if (!msg.SerializeToArray(protoData.data(),
                                      static_cast<int>(protoData.size())))
            {
                lg2::error("[{PATH}]Failed to serialize Arbitration", "PATH",
                           path);
            }
            break;
        }
        case 0x02:
        {
            google::gbmc::nvme::base::PowerManagement msg;
            msg.set_power_state(cdw0 & 0x1F);
            msg.set_workload_hint((cdw0 >> 8) & 0x7);
            protoData.resize(msg.ByteSizeLong());
            if (!msg.SerializeToArray(protoData.data(),
                                      static_cast<int>(protoData.size())))
            {
                lg2::error("[{PATH}]Failed to serialize PowerManagement",
                           "PATH", path);
            }
            break;
        }
        case 0x04:
        {
            google::gbmc::nvme::base::TemperatureThreshold msg;
            msg.set_temperature_threshold(cdw0 & 0xFFFF);
            msg.set_threshold_temperature_select((cdw0 >> 16) & 0xF);
            msg.set_threshold_type_select(
                static_cast<google::gbmc::nvme::base::
                                TemperatureThreshold_ThresholdType>(
                    (cdw0 >> 20) & 0x3));
            protoData.resize(msg.ByteSizeLong());
            if (!msg.SerializeToArray(protoData.data(),
                                      static_cast<int>(protoData.size())))
            {
                lg2::error("[{PATH}]Failed to serialize TemperatureThreshold",
                           "PATH", path);
            }
            break;
        }
        case 0x05:
        {
            google::gbmc::nvme::base::ErrorRecovery msg;
            msg.set_tler(cdw0 & 0xFFFF);
            msg.set_dulbe_enable(
                static_cast<google::gbmc::nvme::base::ErrorRecovery_State>(
                    (cdw0 >> 16) & 0x1));
            protoData.resize(msg.ByteSizeLong());
            if (!msg.SerializeToArray(protoData.data(),
                                      static_cast<int>(protoData.size())))
            {
                lg2::error("[{PATH}]Failed to serialize ErrorRecovery", "PATH",
                           path);
            }
            break;
        }
        case 0x06:
        {
            google::gbmc::nvme::base::VolatileWriteCache msg;
            msg.set_write_cache_enable(
                static_cast<google::gbmc::nvme::base::VolatileWriteCache_State>(
                    cdw0 & 0x1));
            protoData.resize(msg.ByteSizeLong());
            if (!msg.SerializeToArray(protoData.data(),
                                      static_cast<int>(protoData.size())))
            {
                lg2::error("[{PATH}]Failed to serialize VolatileWriteCache",
                           "PATH", path);
            }
            break;
        }
        case 0x07:
        {
            google::gbmc::nvme::base::NumberOfQueues msg;
            msg.set_number_of_submission_queues(cdw0 & 0xFFFF);
            msg.set_number_of_completion_queues((cdw0 >> 16) & 0xFFFF);
            protoData.resize(msg.ByteSizeLong());
            if (!msg.SerializeToArray(protoData.data(),
                                      static_cast<int>(protoData.size())))
            {
                lg2::error("[{PATH}]Failed to serialize NumberOfQueues", "PATH",
                           path);
            }
            break;
        }
        case 0x08:
        {
            google::gbmc::nvme::base::InterruptCoalescing msg;
            msg.set_aggregation_threshold(cdw0 & 0xFF);
            msg.set_aggregation_time((cdw0 >> 8) & 0xFF);
            protoData.resize(msg.ByteSizeLong());
            if (!msg.SerializeToArray(protoData.data(),
                                      static_cast<int>(protoData.size())))
            {
                lg2::error("[{PATH}]Failed to serialize Interrupt Coalescing",
                           "PATH", path);
            }
            break;
        }
        case 0x09:
        {
            google::gbmc::nvme::base::InterruptVectorConfiguration msg;
            msg.set_interrupt_vector(cdw0 & 0xFFFF);
            msg.set_coalescing_disable(
                static_cast<google::gbmc::nvme::base::
                                InterruptVectorConfiguration_CoalescingState>(
                    (cdw0 >> 16) & 0x1));
            protoData.resize(msg.ByteSizeLong());
            if (!msg.SerializeToArray(protoData.data(),
                                      static_cast<int>(protoData.size())))
            {
                lg2::error(
                    "[{PATH}]Failed to serialize InterruptVectorConfiguration",
                    "PATH", path);
            }
            break;
        }
        case 0x0A:
        {
            google::gbmc::nvme::base::WriteAtomicityNormal msg;
            msg.set_disable_normal(
                static_cast<google::gbmc::nvme::base::
                                WriteAtomicityNormal_NormalWriteAtomicityState>(
                    cdw0 & 0x1));
            protoData.resize(msg.ByteSizeLong());
            if (!msg.SerializeToArray(protoData.data(),
                                      static_cast<int>(protoData.size())))
            {
                lg2::error("[{PATH}]Failed to serialize WriteAtomicityNormal",
                           "PATH", path);
            }
            break;
        }
        case 0x0B:
        {
            google::gbmc::nvme::base::AsynchronousEventConfiguration msg;
            msg.set_smart_critical_warning_config(cdw0 & 0xFF);
            protoData.resize(msg.ByteSizeLong());
            if (!msg.SerializeToArray(protoData.data(),
                                      static_cast<int>(protoData.size())))
            {
                lg2::error(
                    "[{PATH}]Failed to serialize AsynchronousEventConfiguration",
                    "PATH", path);
            }
            break;
        }
        case 0xC5:
        {
            if (data.size() < sizeof(OcpLatencyMonitorFeatureBuffer))
            {
                lg2::error(
                    "[{PATH}]OcpLatencyMonitor data buffer too small: {SIZE}",
                    "PATH", path, "SIZE", data.size());
                throw sdbusplus::xyz::openbmc_project::Common::Error::
                    Unavailable();
            }
            OcpLatencyMonitorFeatureBuffer buffer = {};
            std::memcpy(&buffer, data.data(), sizeof(buffer));

            google::gbmc::nvme::base::OcpLatencyMonitor msg;
            msg.set_active_bucket_timer_threshold(
                le16toh(buffer.active_bucket_timer_threshold));
            msg.set_active_threshold_a(buffer.active_threshold_a);
            msg.set_active_threshold_b(buffer.active_threshold_b);
            msg.set_active_threshold_c(buffer.active_threshold_c);
            msg.set_active_threshold_d(buffer.active_threshold_d);
            msg.set_active_latency_config(
                le16toh(buffer.active_latency_config));
            msg.set_active_latency_minimum_window(
                buffer.active_latency_minimum_window);
            msg.set_debug_log_trigger_enable(
                le16toh(buffer.debug_log_trigger_enable));
            msg.set_discard_debug_log(buffer.discard_debug_log);
            msg.set_latency_monitor_feature_enable(
                buffer.latency_monitor_feature_enable);

            protoData.resize(msg.ByteSizeLong());
            if (!msg.SerializeToArray(protoData.data(),
                                      static_cast<int>(protoData.size())))
            {
                lg2::error("[{PATH}]Failed to serialize OcpLatencyMonitor",
                           "PATH", path);
                throw sdbusplus::xyz::openbmc_project::Common::Error::
                    InternalFailure();
            }
            break;
        }
        default:
            lg2::error("[{PATH}]GetFeature: Unsupported FID {FID}", "PATH",
                       path, "FID", lg2::hex, fid);
            throw sdbusplus::xyz::openbmc_project::Common::Error::
                InvalidArgument();
    }

    return protoData;
}

void FeatureStore::setFeature(const boost::asio::yield_context& yield,
                              const std::string& featureName,
                              const std::vector<uint8_t>& data)
{
    std::string path = this->objectPath;
    if (!isFeatureSupported(featureName))
    {
        throw sdbusplus::xyz::openbmc_project::Common::Error::InvalidArgument();
    }
    size_t lastUnderscore = featureName.find_last_of('_');
    if (lastUnderscore == std::string::npos)
    {
        throw sdbusplus::xyz::openbmc_project::Common::Error::InvalidArgument();
    }

    std::string baseName = featureName.substr(0, lastUnderscore);
    std::string selectStr = featureName.substr(lastUnderscore + 1);

    auto it = supportedFeatures.find(baseName);
    if (it == supportedFeatures.end() || it->second.scope != scope)
    {
        throw sdbusplus::xyz::openbmc_project::Common::Error::InvalidArgument();
    }

    if (selectStr == "Default")
    {
        // Default endpoint is read-only
        throw sdbusplus::xyz::openbmc_project::Common::Error::NotAllowed();
    }
    if (selectStr != "Current" && selectStr != "Saved")
    {
        throw sdbusplus::xyz::openbmc_project::Common::Error::InvalidArgument();
    }

    uint8_t fid = it->second.fid;
    bool save = (selectStr == "Saved");

    auto capIt = featureCaps.find(fid);
    if (capIt != featureCaps.end())
    {
        uint32_t caps = capIt->second;
        bool changeable = (caps & nvmeGetFeaturesCapChangeable) != 0;
        bool saveable = (caps & nvmeGetFeaturesCapSaveable) != 0;

        if (save && !saveable)
        {
            lg2::error(
                "[{PATH}]SetFeature: FID {FID} is not saveable per SEL=011b",
                "PATH", path, "FID", lg2::hex, fid);
            throw sdbusplus::xyz::openbmc_project::Common::Error::NotAllowed();
        }

        if (!save && !changeable)
        {
            lg2::error(
                "[{PATH}]SetFeature: FID {FID} is not changeable per SEL=011b",
                "PATH", path, "FID", lg2::hex, fid);
            throw sdbusplus::xyz::openbmc_project::Common::Error::NotAllowed();
        }
    }

    uint32_t cdw11 = 0;
    std::vector<uint8_t> payload;

    switch (fid)
    {
        case 0x01:
        {
            auto msg = parseProto<google::gbmc::nvme::base::Arbitration>(data);
            if (msg.arbitration_burst() > 0xFF ||
                msg.low_priority_weight() > 0xFF ||
                msg.medium_priority_weight() > 0xFF ||
                msg.high_priority_weight() > 0xFF)
            {
                throw sdbusplus::xyz::openbmc_project::Common::Error::
                    InvalidArgument();
            }
            cdw11 = (msg.arbitration_burst() & 0xFF) |
                    ((msg.low_priority_weight() & 0xFF) << 8) |
                    ((msg.medium_priority_weight() & 0xFF) << 16) |
                    ((msg.high_priority_weight() & 0xFF) << 24);
            break;
        }
        case 0x02:
        {
            auto msg =
                parseProto<google::gbmc::nvme::base::PowerManagement>(data);
            if (msg.power_state() > 0x1F || msg.workload_hint() > 0x7)
            {
                throw sdbusplus::xyz::openbmc_project::Common::Error::
                    InvalidArgument();
            }
            cdw11 = (msg.power_state() & 0x1F) |
                    ((msg.workload_hint() & 0x7) << 8);
            break;
        }
        case 0x04:
        {
            auto msg =
                parseProto<google::gbmc::nvme::base::TemperatureThreshold>(
                    data);
            if (msg.temperature_threshold() > 0xFFFF ||
                msg.threshold_temperature_select() > 0xF)
            {
                throw sdbusplus::xyz::openbmc_project::Common::Error::
                    InvalidArgument();
            }
            cdw11 = (msg.temperature_threshold() & 0xFFFF) |
                    ((msg.threshold_temperature_select() & 0xF) << 16) |
                    ((static_cast<uint32_t>(msg.threshold_type_select()) & 0x3)
                     << 20);
            break;
        }
        case 0x05:
        {
            auto msg =
                parseProto<google::gbmc::nvme::base::ErrorRecovery>(data);
            if (msg.tler() > 0xFFFF)
            {
                throw sdbusplus::xyz::openbmc_project::Common::Error::
                    InvalidArgument();
            }
            cdw11 = (msg.tler() & 0xFFFF) |
                    ((static_cast<uint32_t>(msg.dulbe_enable()) & 0x1) << 16);
            break;
        }
        case 0x06:
        {
            auto msg =
                parseProto<google::gbmc::nvme::base::VolatileWriteCache>(data);
            cdw11 = static_cast<uint32_t>(msg.write_cache_enable()) & 0x1;
            break;
        }
        case 0x07:
        {
            auto msg =
                parseProto<google::gbmc::nvme::base::NumberOfQueues>(data);
            if (msg.number_of_submission_queues() > 0xFFFF ||
                msg.number_of_completion_queues() > 0xFFFF)
            {
                throw sdbusplus::xyz::openbmc_project::Common::Error::
                    InvalidArgument();
            }
            cdw11 = (msg.number_of_submission_queues() & 0xFFFF) |
                    ((msg.number_of_completion_queues() & 0xFFFF) << 16);
            break;
        }
        case 0x08:
        {
            auto msg =
                parseProto<google::gbmc::nvme::base::InterruptCoalescing>(data);
            if (msg.aggregation_threshold() > 0xFF ||
                msg.aggregation_time() > 0xFF)
            {
                throw sdbusplus::xyz::openbmc_project::Common::Error::
                    InvalidArgument();
            }
            cdw11 = (msg.aggregation_threshold() & 0xFF) |
                    ((msg.aggregation_time() & 0xFF) << 8);
            break;
        }
        case 0x09:
        {
            auto msg = parseProto<
                google::gbmc::nvme::base::InterruptVectorConfiguration>(data);
            if (msg.interrupt_vector() > 0xFFFF)
            {
                throw sdbusplus::xyz::openbmc_project::Common::Error::
                    InvalidArgument();
            }
            cdw11 = (msg.interrupt_vector() & 0xFFFF) |
                    ((static_cast<uint32_t>(msg.coalescing_disable()) & 0x1)
                     << 16);
            break;
        }
        case 0x0A:
        {
            auto msg =
                parseProto<google::gbmc::nvme::base::WriteAtomicityNormal>(
                    data);
            cdw11 = static_cast<uint32_t>(msg.disable_normal()) & 0x1;
            break;
        }
        case 0x0B:
        {
            auto msg = parseProto<
                google::gbmc::nvme::base::AsynchronousEventConfiguration>(data);
            if (msg.smart_critical_warning_config() > 0xFF)
            {
                throw sdbusplus::xyz::openbmc_project::Common::Error::
                    InvalidArgument();
            }
            cdw11 = msg.smart_critical_warning_config() & 0xFF;
            break;
        }
        case 0xC5:
        {
            auto msg =
                parseProto<google::gbmc::nvme::base::OcpLatencyMonitor>(data);
            if (msg.active_bucket_timer_threshold() > 0xFFFF ||
                msg.active_threshold_a() > 0xFF ||
                msg.active_threshold_b() > 0xFF ||
                msg.active_threshold_c() > 0xFF ||
                msg.active_threshold_d() > 0xFF ||
                msg.active_latency_config() > 0xFFFF ||
                msg.active_latency_minimum_window() > 0xFF ||
                msg.debug_log_trigger_enable() > 0xFFFF ||
                msg.discard_debug_log() > 0xFF ||
                msg.latency_monitor_feature_enable() > 0xFF)
            {
                throw sdbusplus::xyz::openbmc_project::Common::Error::
                    InvalidArgument();
            }
            OcpLatencyMonitorFeatureBuffer buffer = {};
            buffer.active_bucket_timer_threshold = htole16(
                static_cast<uint16_t>(msg.active_bucket_timer_threshold()));
            buffer.active_threshold_a =
                static_cast<uint8_t>(msg.active_threshold_a());
            buffer.active_threshold_b =
                static_cast<uint8_t>(msg.active_threshold_b());
            buffer.active_threshold_c =
                static_cast<uint8_t>(msg.active_threshold_c());
            buffer.active_threshold_d =
                static_cast<uint8_t>(msg.active_threshold_d());
            buffer.active_latency_config =
                htole16(static_cast<uint16_t>(msg.active_latency_config()));
            buffer.active_latency_minimum_window =
                static_cast<uint8_t>(msg.active_latency_minimum_window());
            buffer.debug_log_trigger_enable =
                htole16(static_cast<uint16_t>(msg.debug_log_trigger_enable()));
            buffer.discard_debug_log =
                static_cast<uint8_t>(msg.discard_debug_log());
            buffer.latency_monitor_feature_enable =
                static_cast<uint8_t>(msg.latency_monitor_feature_enable());

            payload.resize(sizeof(buffer), 0);
            std::memcpy(payload.data(), &buffer, sizeof(buffer));
            break;
        }
        default:
            lg2::error("[{PATH}]SetFeature: Unsupported FID {FID}", "PATH",
                       path, "FID", lg2::hex, fid);
            throw sdbusplus::xyz::openbmc_project::Common::Error::
                InvalidArgument();
    }

    NVMeMiIntf::SetFeaturesRequest args{};
    args.fid = fid;
    args.save = save;
    args.cdw11 = cdw11;
    args.nsid = (scope == FeatureScope::Namespace) ? nsid : NVME_NSID_NONE;
    args.data = std::move(payload);

    auto [ex, resp] = asyncAdminSetFeatures(io, nvmeIntf, ctrl, args, yield);

    if (ex)
    {
        lg2::error("[{PATH}]NVMe Set Features (FID {FID}) failed: {ERROR}",
                   "PATH", path, "FID", fid, "ERROR", ex->what());
        throw sdbusplus::xyz::openbmc_project::Common::Error::Unavailable();
    }
}