#include "nvme_features.pb.h"

#include "NVMeError.hpp"
#include "NVMeFeatureStore.hpp"
#include "NVMeMiFake.hpp"

#include <boost/asio/spawn.hpp>
#include <sdbusplus/asio/connection.hpp>
#include <sdbusplus/asio/object_server.hpp>

#include <gmock/gmock.h>
#include <gtest/gtest.h>

class NVMeMiMock :
    public NVMeMiIntf,
    public std::enable_shared_from_this<NVMeMiMock>
{
  public:
    ~NVMeMiMock() override = default;
    MOCK_METHOD(void, start, (const std::shared_ptr<MctpEndpoint>&),
                (override));
    MOCK_METHOD(void, stop, (), (override));
    MOCK_METHOD(void, recover, (), (override));
    MOCK_METHOD(void, miSubsystemHealthStatusPoll,
                (std::function<void(const std::error_code&,
                                    nvme_mi_nvm_ss_health_status*)>&&),
                (override));
    MOCK_METHOD(void, miScanCtrl,
                (std::function<void(const std::error_code&,
                                    const std::vector<nvme_mi_ctrl_t>&)>),
                (override));
    MOCK_METHOD(bool, flushOperations, (std::function<void()>&&), (override));
    MOCK_METHOD(void, adminIdentify,
                (nvme_mi_ctrl_t ctrl, nvme_identify_cns cns, uint32_t nsid,
                 uint16_t cntid,
                 std::function<void(nvme_ex_ptr, std::span<uint8_t>)>&& cb),
                (override));
    MOCK_METHOD(void, adminGetFeatures,
                (nvme_mi_ctrl_t ctrl, const NVMeMiIntf::GetFeaturesRequest& req,
                 std::function<void(nvme_ex_ptr, uint32_t response,
                                    std::span<uint8_t>)>&& cb),
                (override));
    MOCK_METHOD(void, adminSetFeatures,
                (nvme_mi_ctrl_t ctrl, const NVMeMiIntf::SetFeaturesRequest& req,
                 std::function<void(nvme_ex_ptr, uint32_t response)>&& cb),
                (override));
    MOCK_METHOD(
        void, adminGetLogPage,
        (nvme_mi_ctrl_t ctrl, nvme_cmd_get_log_lid lid, uint32_t nsid,
         uint8_t lsp, uint16_t lsi,
         std::function<void(const std::error_code&, std::span<uint8_t>)>&& cb,
         uint8_t dataArea),
        (override));
    MOCK_METHOD(
        void, adminFwCommit,
        (nvme_mi_ctrl_t ctrl, nvme_fw_commit_ca action, uint8_t slot, bool bpid,
         std::function<void(const std::error_code&, nvme_status_field)>&& cb),
        (override));
    MOCK_METHOD(
        void, adminFwDownload,
        (nvme_mi_ctrl_t ctrl, std::string firmwarefile,
         std::function<void(const std::error_code&, nvme_status_field)>&& cb),
        (override));
    MOCK_METHOD(void, adminXfer,
                (nvme_mi_ctrl_t ctrl, const nvme_mi_admin_req_hdr& adminReq,
                 std::span<uint8_t> data, unsigned int timeoutMs,
                 std::function<void(const std::error_code& ec,
                                    const nvme_mi_admin_resp_hdr& adminResp,
                                    std::span<uint8_t> respData)>&& cb),
                (override));
    MOCK_METHOD(void, adminSecuritySend,
                (nvme_mi_ctrl_t ctrl, uint8_t proto, uint16_t specific,
                 std::span<uint8_t> data,
                 std::function<void(const std::error_code&, int)>&& cb),
                (override));
    MOCK_METHOD(void, adminSecurityReceive,
                (nvme_mi_ctrl_t ctrl, uint8_t proto, uint16_t specific,
                 uint32_t al,
                 std::function<void(const std::error_code&, int,
                                    const std::span<uint8_t>)>&& cb),
                (override));
    MOCK_METHOD(
        void, adminNonDataCmd,
        (nvme_mi_ctrl_t ctrl, uint8_t opcode, uint32_t cdw1, uint32_t cdw2,
         uint32_t cdw3, uint32_t cdw10, uint32_t cdw11, uint32_t cdw12,
         uint32_t cdw13, uint32_t cdw14, uint32_t cdw15,
         std::function<void(const std::error_code&, int, uint32_t)>&& cb),
        (override));
    MOCK_METHOD(void, createNamespace,
                (nvme_mi_ctrl_t ctrl, uint64_t size, size_t lbaFormat,
                 bool metadataAtEnd,
                 std::function<void(nvme_ex_ptr)>&& submittedCb,
                 std::function<void(nvme_ex_ptr, NVMeNSIdentify)>&& finishedCb),
                (override));
    MOCK_METHOD(void, adminDeleteNamespace,
                (nvme_mi_ctrl_t ctrl, uint32_t nsid,
                 std::function<void(const std::error_code&, int)>&& cb),
                (override));
    MOCK_METHOD(void, adminListNamespaces,
                (nvme_mi_ctrl_t ctrl,
                 std::function<void(nvme_ex_ptr, std::vector<uint32_t>)>&& cb),
                (override));
    MOCK_METHOD(void, adminAttachDetachNamespace,
                (nvme_mi_ctrl_t ctrl, uint16_t ctrlid, uint32_t nsid,
                 bool attach,
                 std::function<void(const std::error_code&, int)>&& cb),
                (override));
    MOCK_METHOD(void, adminSanitize,
                (nvme_mi_ctrl_t ctrl, enum nvme_sanitize_sanact sanact,
                 uint8_t passes, uint32_t pattern, bool invertPattern,
                 std::function<void(nvme_ex_ptr)>&& cb),
                (override));
    MOCK_METHOD(
        void, adminLockdownInband,
        (nvme_mi_ctrl_t ctrl, uint8_t prohibit,
         const std::vector<uint8_t>& adminCmds,
         const std::vector<uint8_t>& features,
         const std::vector<uint8_t>& logPages,
         std::function<
             void(const std::error_code&, int nvmeStatus, uint32_t comptionDw0,
                  const std::string& failureScope, uint32_t failureId)>&& cb),
        (override));
    MOCK_METHOD(
        void, adminFormat,
        (nvme_mi_ctrl_t ctrl, uint32_t nsid, uint8_t lbaFormat,
         uint8_t secureEraseSetting, uint8_t metadataSettings,
         uint8_t dataProtectionType, uint8_t protectionInformationLocation,
         std::function<void(const std::error_code&, int nvmeStatus)>&& cb),
        (override));
};

class FeatureStoreTest : public ::testing::Test
{
  public:
    static std::vector<uint8_t>
        callGetFeature(const std::shared_ptr<FeatureStore>& fs,
                       const boost::asio::yield_context& yield,
                       const std::string& featureName)
    {
        return fs->getFeature(yield, featureName);
    }

    static void callSetFeature(const std::shared_ptr<FeatureStore>& fs,
                               const boost::asio::yield_context& yield,
                               const std::string& featureName,
                               const std::vector<uint8_t>& data)
    {
        fs->setFeature(yield, featureName, data);
    }

    static std::map<uint8_t, uint32_t>&
        getFeatureCaps(const std::shared_ptr<FeatureStore>& fs)
    {
        return fs->featureCaps;
    }

    static bool callIsFeatureSupported(const std::shared_ptr<FeatureStore>& fs,
                                       const std::string& featureName)
    {
        return fs->isFeatureSupported(featureName);
    }

    static void addSupportedFeature(const std::shared_ptr<FeatureStore>& fs,
                                    const std::string& name,
                                    const std::string& fid,
                                    const std::string& protoClass)
    {
        fs->featureCollection.emplace_back(name, fid, protoClass);
    }
};

TEST_F(FeatureStoreTest, GetFeatureArbitration)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/ctrl0";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 0,
        FeatureScope::Controller);

    FeatureStoreTest::addSupportedFeature(featureStore, "Arbitration_Current",
                                          "0x01",
                                          "google.gbmc.nvme.base.Arbitration");

    EXPECT_CALL(*mockNvme,
                adminGetFeatures(::testing::_, ::testing::_, ::testing::_))
        .WillOnce(
            [](nvme_mi_ctrl_t, const NVMeMiIntf::GetFeaturesRequest& req,
               std::function<void(nvme_ex_ptr, uint32_t, std::span<uint8_t>)>&&
                   cb) {
        EXPECT_EQ(req.fid, 0x01);    // Arbitration
        cb(nullptr, 0x01020304, {}); // burst=4, low=3, med=2, high=1
    });

    std::vector<uint8_t> result;
    bool done = false;
    boost::asio::spawn(io, [&](const boost::asio::yield_context& yield) {
        result = FeatureStoreTest::callGetFeature(featureStore, yield,
                                                  "Arbitration_Current");
        done = true;
    });
    while (!done)
    {
        io.run_one();
    }
    EXPECT_TRUE(done);

    google::gbmc::nvme::base::Arbitration msg;
    ASSERT_TRUE(
        msg.ParseFromArray(result.data(), static_cast<int>(result.size())));
    EXPECT_EQ(msg.arbitration_burst(), 0x04);
    EXPECT_EQ(msg.low_priority_weight(), 0x03);
    EXPECT_EQ(msg.medium_priority_weight(), 0x02);
    EXPECT_EQ(msg.high_priority_weight(), 0x01);
}

TEST_F(FeatureStoreTest, SetFeatureArbitration)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/ctrl0";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 0,
        FeatureScope::Controller);

    FeatureStoreTest::addSupportedFeature(featureStore, "Arbitration_Current",
                                          "0x01",
                                          "google.gbmc.nvme.base.Arbitration");

    google::gbmc::nvme::base::Arbitration msg;
    msg.set_arbitration_burst(0x08);
    msg.set_low_priority_weight(0x07);
    msg.set_medium_priority_weight(0x06);
    msg.set_high_priority_weight(0x05);

    std::vector<uint8_t> data(msg.ByteSizeLong());
    ASSERT_TRUE(
        msg.SerializeToArray(data.data(), static_cast<int>(data.size())));

    EXPECT_CALL(*mockNvme,
                adminSetFeatures(::testing::_, ::testing::_, ::testing::_))
        .WillOnce([](nvme_mi_ctrl_t, const NVMeMiIntf::SetFeaturesRequest& req,
                     std::function<void(nvme_ex_ptr, uint32_t)>&& cb) {
        EXPECT_EQ(req.fid, 0x01); // Arbitration
        EXPECT_EQ(req.cdw11, 0x05060708);
        cb(nullptr, 0);
    });

    bool done = false;
    boost::asio::spawn(io, [&](const boost::asio::yield_context& yield) {
        FeatureStoreTest::callSetFeature(featureStore, yield,
                                         "Arbitration_Current", data);
        done = true;
    });
    while (!done)
    {
        io.run_one();
    }
    EXPECT_TRUE(done);
}

TEST_F(FeatureStoreTest, DiscoverFeaturesONCS)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/ctrl0";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 0,
        FeatureScope::Controller);

    // Mock Identify Controller to return ONCS with bit 4 set
    EXPECT_CALL(*mockNvme,
                adminIdentify(::testing::_, NVME_IDENTIFY_CNS_CTRL,
                              ::testing::_, ::testing::_, ::testing::_))
        .WillOnce(
            [](nvme_mi_ctrl_t, nvme_identify_cns, uint32_t, uint16_t,
               std::function<void(nvme_ex_ptr, std::span<uint8_t>)>&& cb) {
        nvme_id_ctrl id{};
        id.vwc = 1;                // VWC supported
        id.oncs = htole16(1 << 4); // Save and Select supported
        std::vector<uint8_t> data(sizeof(id));
        memcpy(data.data(), &id, sizeof(id));
        cb(nullptr, data);
    });

    // Mock Get Features for capabilities of Arbitration (FID 0x01)
    // Return Saveable (bit 2) and Changeable (bit 0)
    EXPECT_CALL(
        *mockNvme,
        adminGetFeatures(::testing::_,
                         ::testing::Field(&NVMeMiIntf::GetFeaturesRequest::sel,
                                          NVME_GET_FEATURES_SEL_SUPPORTED),
                         ::testing::_))
        .WillRepeatedly(
            [](nvme_mi_ctrl_t, const NVMeMiIntf::GetFeaturesRequest& req,
               std::function<void(nvme_ex_ptr, uint32_t, std::span<uint8_t>)>&&
                   cb) {
        if (req.fid == 0x01)
        {
            cb(nullptr, 0x5, {}); // Changeable | Saveable
        }
        else
        {
            cb(nullptr, 0x4, {}); // Just Changeable
        }
    });

    // Mock Get Log Page for Log Page 12h
    EXPECT_CALL(*mockNvme, adminGetLogPage(
                               ::testing::_, NVME_LOG_LID_FID_SUPPORTED_EFFECTS,
                               ::testing::_, ::testing::_, ::testing::_,
                               ::testing::_, ::testing::_))
        .WillOnce(
            [](nvme_mi_ctrl_t, nvme_cmd_get_log_lid, uint32_t, uint8_t,
               uint16_t,
               std::function<void(const std::error_code&, std::span<uint8_t>)>&&
                   cb,
               uint8_t) {
        std::vector<uint8_t> data(sizeof(struct nvme_fid_supported_effects_log),
                                  0);
        std::vector<uint8_t> fids = {0x01, 0x02, 0x04, 0x05, 0x06,
                                     0x07, 0x08, 0x09, 0x0A, 0x0B};
        for (uint8_t fid : fids)
        {
            data[static_cast<size_t>(fid) * 4] = 0x1;
        }
        cb(std::error_code(), data);
    });

    bool done = false;
    featureStore->init([&done](const std::error_code& ec) {
        EXPECT_FALSE(ec);
        done = true;
    });

    while (!done)
    {
        io.run_one();
    }

    auto caps = FeatureStoreTest::getFeatureCaps(featureStore);

    // Check for Arbitration states (FID 0x01)
    auto it = caps.find(0x01);
    ASSERT_TRUE(it != caps.end());
    uint32_t cap = it->second;

    EXPECT_TRUE(cap & 0x1); // Saveable
    EXPECT_TRUE(cap & 0x4); // Changeable
}

TEST_F(FeatureStoreTest, DiscoverFeaturesCapabilityErrorSkipsUnsupportedFeature)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/vol1";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 1,
        FeatureScope::Namespace);

    // Mock Identify Controller to return ONCS with bit 4 set (Save/Select
    // supported)
    EXPECT_CALL(*mockNvme,
                adminIdentify(::testing::_, NVME_IDENTIFY_CNS_CTRL,
                              ::testing::_, ::testing::_, ::testing::_))
        .WillOnce(
            [](nvme_mi_ctrl_t, nvme_identify_cns, uint32_t, uint16_t,
               std::function<void(nvme_ex_ptr, std::span<uint8_t>)>&& cb) {
        nvme_id_ctrl id{};
        id.vwc = 1;                // VWC supported
        id.oncs = htole16(1 << 4); // Save and Select supported
        std::vector<uint8_t> data(sizeof(id));
        memcpy(data.data(), &id, sizeof(id));
        cb(nullptr, data);
    });

    // Mock Get Log Page to fail or return empty
    EXPECT_CALL(*mockNvme, adminGetLogPage(
                               ::testing::_, NVME_LOG_LID_FID_SUPPORTED_EFFECTS,
                               ::testing::_, ::testing::_, ::testing::_,
                               ::testing::_, ::testing::_))
        .WillOnce(
            [](nvme_mi_ctrl_t, nvme_cmd_get_log_lid, uint32_t, uint8_t,
               uint16_t,
               std::function<void(const std::error_code&, std::span<uint8_t>)>&&
                   cb,
               uint8_t) {
        cb(std::make_error_code(std::errc::bad_message), {});
    });

    // Mock Get Features SEL=011b (capabilities query)
    // ErrorRecovery (FID 0x05) fails with Invalid Field in Command (0x02)
    EXPECT_CALL(
        *mockNvme,
        adminGetFeatures(::testing::_,
                         ::testing::Field(&NVMeMiIntf::GetFeaturesRequest::sel,
                                          NVME_GET_FEATURES_SEL_SUPPORTED),
                         ::testing::_))
        .WillRepeatedly(
            [](nvme_mi_ctrl_t, const NVMeMiIntf::GetFeaturesRequest& req,
               std::function<void(nvme_ex_ptr, uint32_t, std::span<uint8_t>)>&&
                   cb) {
        if (req.fid == 0x05)
        {
            // Simulate drive returning Invalid Field in Command
            cb(makeLibNVMeError(0, 0x02, "Invalid Field in Command"), 0, {});
        }
        else
        {
            cb(nullptr, 0x4, {}); // Changeable
        }
    });

    bool done = false;
    featureStore->init([&done](const std::error_code& ec) {
        EXPECT_EQ(ec, std::errc::not_supported);
        done = true;
    });

    while (!done)
    {
        io.run_one();
    }

    // Verify ErrorRecovery_Current is NOT in featureStore
    EXPECT_FALSE(FeatureStoreTest::callIsFeatureSupported(
        featureStore, "ErrorRecovery_Current"));
}

TEST_F(FeatureStoreTest, GetFeatureInvalidName)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/ctrl0";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 0,
        FeatureScope::Controller);

    bool threw = false;
    bool done = false;
    boost::asio::spawn(io, [&](const boost::asio::yield_context& yield) {
        try
        {
            FeatureStoreTest::callGetFeature(featureStore, yield,
                                             "InvalidFeature_Current");
        }
        catch (const sdbusplus::xyz::openbmc_project::Common::Error::
                   InvalidArgument&)
        {
            threw = true;
        }
        done = true;
    });
    while (!done)
    {
        io.run_one();
    }
    EXPECT_TRUE(threw);
}

TEST_F(FeatureStoreTest, GetFeatureDbusError)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/ctrl0";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 0,
        FeatureScope::Controller);

    FeatureStoreTest::addSupportedFeature(featureStore, "Arbitration_Current",
                                          "0x01",
                                          "google.gbmc.nvme.base.Arbitration");

    EXPECT_CALL(*mockNvme,
                adminGetFeatures(::testing::_, ::testing::_, ::testing::_))
        .WillOnce(
            [](nvme_mi_ctrl_t, const NVMeMiIntf::GetFeaturesRequest&,
               std::function<void(nvme_ex_ptr, uint32_t, std::span<uint8_t>)>&&
                   cb) {
        cb(makeLibNVMeError(std::make_error_code(std::errc::connection_aborted),
                            -1, "mock"),
           0, {});
    });

    bool threw = false;
    bool done = false;
    boost::asio::spawn(io, [&](const boost::asio::yield_context& yield) {
        try
        {
            FeatureStoreTest::callGetFeature(featureStore, yield,
                                             "Arbitration_Current");
        }
        catch (
            const sdbusplus::xyz::openbmc_project::Common::Error::Unavailable&)
        {
            threw = true;
        }
        done = true;
    });
    while (!done)
    {
        io.run_one();
    }
    EXPECT_TRUE(threw);
}

TEST_F(FeatureStoreTest, SetFeatureInvalidName)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/ctrl0";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 0,
        FeatureScope::Controller);

    bool threw = false;
    bool done = false;
    boost::asio::spawn(io, [&](const boost::asio::yield_context& yield) {
        try
        {
            FeatureStoreTest::callSetFeature(featureStore, yield,
                                             "Arbitration_InvalidSelect", {});
        }
        catch (const sdbusplus::xyz::openbmc_project::Common::Error::
                   InvalidArgument&)
        {
            threw = true;
        }
        done = true;
    });
    while (!done)
    {
        io.run_one();
    }
    EXPECT_TRUE(threw);
}

TEST_F(FeatureStoreTest, SetFeatureProtoDeserializationFailure)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/ctrl0";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 0,
        FeatureScope::Controller);

    FeatureStoreTest::addSupportedFeature(featureStore, "Arbitration_Current",
                                          "0x01",
                                          "google.gbmc.nvme.base.Arbitration");

    std::vector<uint8_t> malformedData = {0xFF, 0xFF, 0xFF};

    bool threw = false;
    bool done = false;
    boost::asio::spawn(io, [&](const boost::asio::yield_context& yield) {
        try
        {
            FeatureStoreTest::callSetFeature(
                featureStore, yield, "Arbitration_Current", malformedData);
        }
        catch (const sdbusplus::xyz::openbmc_project::Common::Error::
                   InvalidArgument&)
        {
            threw = true;
        }
        done = true;
    });
    while (!done)
    {
        io.run_one();
    }
    EXPECT_TRUE(threw);
}

TEST_F(FeatureStoreTest, SetFeatureInvalidArgumentBounds)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/ctrl0";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 0,
        FeatureScope::Controller);

    FeatureStoreTest::addSupportedFeature(featureStore, "Arbitration_Current",
                                          "0x01",
                                          "google.gbmc.nvme.base.Arbitration");

    google::gbmc::nvme::base::Arbitration msg;
    msg.set_arbitration_burst(256); // Exceeds 8-bit bound (0-255)

    std::vector<uint8_t> data(msg.ByteSizeLong());
    ASSERT_TRUE(
        msg.SerializeToArray(data.data(), static_cast<int>(data.size())));

    bool threw = false;
    bool done = false;
    boost::asio::spawn(io, [&](const boost::asio::yield_context& yield) {
        try
        {
            FeatureStoreTest::callSetFeature(featureStore, yield,
                                             "Arbitration_Current", data);
        }
        catch (const sdbusplus::xyz::openbmc_project::Common::Error::
                   InvalidArgument&)
        {
            threw = true;
        }
        done = true;
    });
    while (!done)
    {
        io.run_one();
    }
    EXPECT_TRUE(threw);
}

TEST_F(FeatureStoreTest, SetFeatureNotChangeableOrSaveable)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/ctrl0";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 0,
        FeatureScope::Controller);

    FeatureStoreTest::addSupportedFeature(featureStore, "Arbitration_Current",
                                          "0x01",
                                          "google.gbmc.nvme.base.Arbitration");
    FeatureStoreTest::addSupportedFeature(featureStore, "Arbitration_Saved",
                                          "0x01",
                                          "google.gbmc.nvme.base.Arbitration");

    // Seed featureCaps for Arbitration (FID 0x01) with 0 (neither changeable
    // nor saveable)
    FeatureStoreTest::getFeatureCaps(featureStore)[0x01] = 0;

    google::gbmc::nvme::base::Arbitration msg;
    msg.set_arbitration_burst(4);
    std::vector<uint8_t> data(msg.ByteSizeLong());
    ASSERT_TRUE(
        msg.SerializeToArray(data.data(), static_cast<int>(data.size())));

    // 1. Try to set Current (changeable=false) -> should throw NotAllowed
    bool threwNotAllowedCurrent = false;
    bool doneCurrent = false;
    boost::asio::spawn(io, [&](const boost::asio::yield_context& yield) {
        try
        {
            FeatureStoreTest::callSetFeature(featureStore, yield,
                                             "Arbitration_Current", data);
        }
        catch (
            const sdbusplus::xyz::openbmc_project::Common::Error::NotAllowed&)
        {
            threwNotAllowedCurrent = true;
        }
        doneCurrent = true;
    });
    while (!doneCurrent)
    {
        io.run_one();
    }
    EXPECT_TRUE(threwNotAllowedCurrent);

    // 2. Try to set Saved (saveable=false) -> should throw NotAllowed
    io.restart();
    bool threwNotAllowedSaved = false;
    bool doneSaved = false;
    boost::asio::spawn(io, [&](const boost::asio::yield_context& yield) {
        try
        {
            FeatureStoreTest::callSetFeature(featureStore, yield,
                                             "Arbitration_Saved", data);
        }
        catch (
            const sdbusplus::xyz::openbmc_project::Common::Error::NotAllowed&)
        {
            threwNotAllowedSaved = true;
        }
        doneSaved = true;
    });
    while (!doneSaved)
    {
        io.run_one();
    }
    EXPECT_TRUE(threwNotAllowedSaved);
}

TEST_F(FeatureStoreTest, DiscoverFeaturesIdentifyControllerFails)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/vol1";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 1,
        FeatureScope::Namespace);

    // Mock Identify Controller to fail (e.g. absent drive)
    EXPECT_CALL(*mockNvme,
                adminIdentify(::testing::_, NVME_IDENTIFY_CNS_CTRL,
                              ::testing::_, ::testing::_, ::testing::_))
        .WillOnce(
            [](nvme_mi_ctrl_t, nvme_identify_cns, uint32_t, uint16_t,
               std::function<void(nvme_ex_ptr, std::span<uint8_t>)>&& cb) {
        cb(makeLibNVMeError(0, 0x02, "Identify Controller failed"), {});
    });

    bool done = false;
    std::error_code resultEc;
    featureStore->init([&done, &resultEc](const std::error_code& ec) {
        resultEc = ec;
        done = true;
    });

    while (!done)
    {
        io.run_one();
    }

    EXPECT_TRUE(done);
    EXPECT_EQ(resultEc, std::make_error_code(std::errc::no_such_device));
    EXPECT_FALSE(FeatureStoreTest::callIsFeatureSupported(
        featureStore, "TemperatureThreshold_Current"));
}
TEST_F(FeatureStoreTest, DiscoverFeaturesNoONCSBit4Supported)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/ctrl0";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 0,
        FeatureScope::Controller);

    // Mock Identify Controller to return ONCS without bit 4
    EXPECT_CALL(*mockNvme,
                adminIdentify(::testing::_, NVME_IDENTIFY_CNS_CTRL,
                              ::testing::_, ::testing::_, ::testing::_))
        .WillOnce(
            [](nvme_mi_ctrl_t, nvme_identify_cns, uint32_t, uint16_t,
               std::function<void(nvme_ex_ptr, std::span<uint8_t>)>&& cb) {
        nvme_id_ctrl id{};
        id.vwc = 1;
        id.oncs = 0; // Bit 4 NOT set
        std::vector<uint8_t> data(sizeof(id));
        memcpy(data.data(), &id, sizeof(id));
        cb(nullptr, data);
    });

    // Mock Log Page 12h to fail
    EXPECT_CALL(*mockNvme, adminGetLogPage(
                               ::testing::_, NVME_LOG_LID_FID_SUPPORTED_EFFECTS,
                               ::testing::_, ::testing::_, ::testing::_,
                               ::testing::_, ::testing::_))
        .WillOnce(
            [](nvme_mi_ctrl_t, nvme_cmd_get_log_lid, uint32_t, uint8_t,
               uint16_t,
               std::function<void(const std::error_code&, std::span<uint8_t>)>&&
                   cb,
               uint8_t) {
        cb(std::make_error_code(std::errc::bad_message), {});
    });

    // Mock Get Features with SEL=CURRENT (000b)
    // Only Arbitration (FID 0x01) succeeds; others fail
    EXPECT_CALL(
        *mockNvme,
        adminGetFeatures(::testing::_,
                         ::testing::Field(&NVMeMiIntf::GetFeaturesRequest::sel,
                                          NVME_GET_FEATURES_SEL_CURRENT),
                         ::testing::_))
        .WillRepeatedly(
            [](nvme_mi_ctrl_t, const NVMeMiIntf::GetFeaturesRequest& req,
               std::function<void(nvme_ex_ptr, uint32_t, std::span<uint8_t>)>&&
                   cb) {
        if (req.fid == 0x01)
        {
            cb(nullptr, 0, {});
        }
        else
        {
            cb(makeLibNVMeError(0, 0x02, "Feature not supported"), 0, {});
        }
    });

    bool done = false;
    featureStore->init([&done](const std::error_code& ec) {
        EXPECT_FALSE(ec);
        done = true;
    });

    while (!done)
    {
        io.run_one();
    }

    EXPECT_TRUE(FeatureStoreTest::callIsFeatureSupported(
        featureStore, "Arbitration_Current"));
    EXPECT_FALSE(FeatureStoreTest::callIsFeatureSupported(
        featureStore, "PowerManagement_Current"));
    EXPECT_FALSE(FeatureStoreTest::callIsFeatureSupported(
        featureStore, "Arbitration_Default"));
    EXPECT_FALSE(FeatureStoreTest::callIsFeatureSupported(featureStore,
                                                          "Arbitration_Saved"));
}

TEST_F(FeatureStoreTest, DiscoverFeaturesNoONCSBit4AllFeaturesFail)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/ctrl0";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 0,
        FeatureScope::Controller);

    // Mock Identify Controller to return ONCS without bit 4
    EXPECT_CALL(*mockNvme,
                adminIdentify(::testing::_, NVME_IDENTIFY_CNS_CTRL,
                              ::testing::_, ::testing::_, ::testing::_))
        .WillOnce(
            [](nvme_mi_ctrl_t, nvme_identify_cns, uint32_t, uint16_t,
               std::function<void(nvme_ex_ptr, std::span<uint8_t>)>&& cb) {
        nvme_id_ctrl id{};
        id.vwc = 1;
        id.oncs = 0; // Bit 4 NOT set
        std::vector<uint8_t> data(sizeof(id));
        memcpy(data.data(), &id, sizeof(id));
        cb(nullptr, data);
    });

    // Mock Log Page 12h to fail
    EXPECT_CALL(*mockNvme, adminGetLogPage(
                               ::testing::_, NVME_LOG_LID_FID_SUPPORTED_EFFECTS,
                               ::testing::_, ::testing::_, ::testing::_,
                               ::testing::_, ::testing::_))
        .WillOnce(
            [](nvme_mi_ctrl_t, nvme_cmd_get_log_lid, uint32_t, uint8_t,
               uint16_t,
               std::function<void(const std::error_code&, std::span<uint8_t>)>&&
                   cb,
               uint8_t) {
        cb(std::make_error_code(std::errc::bad_message), {});
    });

    // Mock Get Features with SEL=CURRENT (000b) to fail for all features (e.g.
    // IronFist)
    EXPECT_CALL(
        *mockNvme,
        adminGetFeatures(::testing::_,
                         ::testing::Field(&NVMeMiIntf::GetFeaturesRequest::sel,
                                          NVME_GET_FEATURES_SEL_CURRENT),
                         ::testing::_))
        .WillRepeatedly(
            [](nvme_mi_ctrl_t, const NVMeMiIntf::GetFeaturesRequest&,
               std::function<void(nvme_ex_ptr, uint32_t, std::span<uint8_t>)>&&
                   cb) {
        cb(makeLibNVMeError(0, 0xf3, "Vendor Specific Status (SCT 7 SC 0xf3)"),
           0, {});
    });

    bool done = false;
    std::error_code resultEc;
    featureStore->init([&done, &resultEc](const std::error_code& ec) {
        resultEc = ec;
        done = true;
    });

    while (!done)
    {
        io.run_one();
    }

    EXPECT_TRUE(done);
    EXPECT_EQ(resultEc, std::make_error_code(std::errc::not_supported));
    EXPECT_FALSE(FeatureStoreTest::callIsFeatureSupported(
        featureStore, "Arbitration_Current"));
}

TEST_F(FeatureStoreTest, GetFeatureOcpLatencyMonitor)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/ctrl0";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 0,
        FeatureScope::Controller);

    FeatureStoreTest::addSupportedFeature(
        featureStore, "OcpLatencyMonitor_Current", "0xc5",
        "google.gbmc.nvme.base.OcpLatencyMonitor");

    std::vector<uint8_t> mockData(4096, 0);
    // Byte 0-1: 1000 ms (0x03E8)
    mockData[0] = 0xE8;
    mockData[1] = 0x03;
    mockData[2] = 1; // Threshold A
    mockData[3] = 2; // Threshold B
    mockData[4] = 3; // Threshold C
    mockData[5] = 4; // Threshold D
    // Byte 6-7: 5 (0x0005)
    mockData[6] = 5;
    mockData[7] = 0;
    mockData[8] = 6; // Minimum Window
    // Byte 9-10: 7 (0x0007)
    mockData[9] = 7;
    mockData[10] = 0;
    mockData[11] = 1; // Discard Debug Log
    mockData[12] = 1; // Feature Enable

    EXPECT_CALL(*mockNvme,
                adminGetFeatures(::testing::_, ::testing::_, ::testing::_))
        .WillOnce([&mockData](nvme_mi_ctrl_t,
                              const NVMeMiIntf::GetFeaturesRequest& req,
                              std::function<void(nvme_ex_ptr, uint32_t,
                                                 std::span<uint8_t>)>&& cb) {
        EXPECT_EQ(req.fid, 0xC5); // OcpLatencyMonitor
        EXPECT_EQ(req.data_len, 4096);
        cb(nullptr, 0, mockData);
    });

    std::vector<uint8_t> result;
    bool done = false;
    boost::asio::spawn(io, [&](const boost::asio::yield_context& yield) {
        result = FeatureStoreTest::callGetFeature(featureStore, yield,
                                                  "OcpLatencyMonitor_Current");
        done = true;
    });
    while (!done)
    {
        io.run_one();
    }
    EXPECT_TRUE(done);

    google::gbmc::nvme::base::OcpLatencyMonitor msg;
    ASSERT_TRUE(
        msg.ParseFromArray(result.data(), static_cast<int>(result.size())));
    EXPECT_EQ(msg.active_bucket_timer_threshold(), 1000);
    EXPECT_EQ(msg.active_threshold_a(), 1);
    EXPECT_EQ(msg.active_threshold_b(), 2);
    EXPECT_EQ(msg.active_threshold_c(), 3);
    EXPECT_EQ(msg.active_threshold_d(), 4);
    EXPECT_EQ(msg.active_latency_config(), 5);
    EXPECT_EQ(msg.active_latency_minimum_window(), 6);
    EXPECT_EQ(msg.debug_log_trigger_enable(), 7);
    EXPECT_EQ(msg.discard_debug_log(), 1);
    EXPECT_EQ(msg.latency_monitor_feature_enable(), 1);
}

TEST_F(FeatureStoreTest, GetFeatureOcpLatencyMonitorBufferTooSmall)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/ctrl0";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 0,
        FeatureScope::Controller);

    FeatureStoreTest::addSupportedFeature(
        featureStore, "OcpLatencyMonitor_Current", "0xc5",
        "google.gbmc.nvme.base.OcpLatencyMonitor");

    std::vector<uint8_t> mockData(10, 0); // Less than 13 bytes

    EXPECT_CALL(*mockNvme,
                adminGetFeatures(::testing::_, ::testing::_, ::testing::_))
        .WillOnce([&mockData](nvme_mi_ctrl_t,
                              const NVMeMiIntf::GetFeaturesRequest& req,
                              std::function<void(nvme_ex_ptr, uint32_t,
                                                 std::span<uint8_t>)>&& cb) {
        EXPECT_EQ(req.fid, 0xC5); // OcpLatencyMonitor
        EXPECT_EQ(req.data_len, 4096);
        cb(nullptr, 0, mockData);
    });

    bool threw = false;
    bool done = false;
    boost::asio::spawn(io, [&](const boost::asio::yield_context& yield) {
        try
        {
            FeatureStoreTest::callGetFeature(featureStore, yield,
                                             "OcpLatencyMonitor_Current");
        }
        catch (
            const sdbusplus::xyz::openbmc_project::Common::Error::Unavailable&)
        {
            threw = true;
        }
        done = true;
    });
    while (!done)
    {
        io.run_one();
    }
    EXPECT_TRUE(done);
    EXPECT_TRUE(threw);
}

TEST_F(FeatureStoreTest, SetFeatureOcpLatencyMonitor)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/ctrl0";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 0,
        FeatureScope::Controller);

    FeatureStoreTest::addSupportedFeature(
        featureStore, "OcpLatencyMonitor_Current", "0xc5",
        "google.gbmc.nvme.base.OcpLatencyMonitor");

    google::gbmc::nvme::base::OcpLatencyMonitor msg;
    msg.set_active_bucket_timer_threshold(1000);
    msg.set_active_threshold_a(1);
    msg.set_active_threshold_b(2);
    msg.set_active_threshold_c(3);
    msg.set_active_threshold_d(4);
    msg.set_active_latency_config(5);
    msg.set_active_latency_minimum_window(6);
    msg.set_debug_log_trigger_enable(7);
    msg.set_discard_debug_log(1);
    msg.set_latency_monitor_feature_enable(1);

    std::vector<uint8_t> data(msg.ByteSizeLong());
    ASSERT_TRUE(
        msg.SerializeToArray(data.data(), static_cast<int>(data.size())));

    EXPECT_CALL(*mockNvme,
                adminSetFeatures(::testing::_, ::testing::_, ::testing::_))
        .WillOnce([](nvme_mi_ctrl_t, const NVMeMiIntf::SetFeaturesRequest& req,
                     std::function<void(nvme_ex_ptr, uint32_t)>&& cb) {
        EXPECT_EQ(req.fid, 0xC5); // OcpLatencyMonitor
        EXPECT_EQ(req.data.size(), 4096);
        EXPECT_EQ(req.data[0], 0xE8);
        EXPECT_EQ(req.data[1], 0x03);
        EXPECT_EQ(req.data[2], 1);
        EXPECT_EQ(req.data[3], 2);
        EXPECT_EQ(req.data[4], 3);
        EXPECT_EQ(req.data[5], 4);
        EXPECT_EQ(req.data[6], 5);
        EXPECT_EQ(req.data[7], 0);
        EXPECT_EQ(req.data[8], 6);
        EXPECT_EQ(req.data[9], 7);
        EXPECT_EQ(req.data[10], 0);
        EXPECT_EQ(req.data[11], 1);
        EXPECT_EQ(req.data[12], 1);
        for (size_t i = 13; i < 4096; ++i)
        {
            EXPECT_EQ(req.data[i], 0);
        }
        cb(nullptr, 0);
    });

    bool done = false;
    boost::asio::spawn(io, [&](const boost::asio::yield_context& yield) {
        FeatureStoreTest::callSetFeature(featureStore, yield,
                                         "OcpLatencyMonitor_Current", data);
        done = true;
    });
    while (!done)
    {
        io.run_one();
    }
    EXPECT_TRUE(done);
}

TEST_F(FeatureStoreTest, SetFeatureOcpLatencyMonitorInvalidArgumentBounds)
{
    boost::asio::io_context io;
    auto conn = std::make_shared<sdbusplus::asio::connection>(io);
    sdbusplus::asio::object_server objServer(conn);
    std::string objectPath = "/xyz/openbmc_project/nvme/ctrl0";

    auto mockNvme = std::make_shared<NVMeMiMock>();
    nvme_mi_ctrl_t ctrl = nullptr;

    auto featureStore = std::make_shared<FeatureStore>(
        io, objServer, conn, objectPath, mockNvme, ctrl, 0,
        FeatureScope::Controller);

    FeatureStoreTest::addSupportedFeature(
        featureStore, "OcpLatencyMonitor_Current", "0xc5",
        "google.gbmc.nvme.base.OcpLatencyMonitor");

    google::gbmc::nvme::base::OcpLatencyMonitor msg;
    // Set field out of uint16 bounds
    msg.set_active_bucket_timer_threshold(0x10000);

    std::vector<uint8_t> data(msg.ByteSizeLong());
    ASSERT_TRUE(
        msg.SerializeToArray(data.data(), static_cast<int>(data.size())));

    bool threw = false;
    bool done = false;
    boost::asio::spawn(io, [&](const boost::asio::yield_context& yield) {
        try
        {
            FeatureStoreTest::callSetFeature(featureStore, yield,
                                             "OcpLatencyMonitor_Current", data);
        }
        catch (const sdbusplus::xyz::openbmc_project::Common::Error::
                   InvalidArgument&)
        {
            threw = true;
        }
        done = true;
    });
    while (!done)
    {
        io.run_one();
    }
    EXPECT_TRUE(threw);
}

TEST_F(FeatureStoreTest, GetFeatureDataLen)
{
    EXPECT_EQ(FeatureStore::getFeatureDataLen(0x01), 0);
    EXPECT_EQ(FeatureStore::getFeatureDataLen(0x02), 0);
    EXPECT_EQ(FeatureStore::getFeatureDataLen(0x04), 0);
    EXPECT_EQ(FeatureStore::getFeatureDataLen(0x05), 0);
    EXPECT_EQ(FeatureStore::getFeatureDataLen(0x06), 0);
    EXPECT_EQ(FeatureStore::getFeatureDataLen(0x07), 0);
    EXPECT_EQ(FeatureStore::getFeatureDataLen(0x08), 0);
    EXPECT_EQ(FeatureStore::getFeatureDataLen(0x09), 0);
    EXPECT_EQ(FeatureStore::getFeatureDataLen(0x0A), 0);
    EXPECT_EQ(FeatureStore::getFeatureDataLen(0x0B), 0);
    EXPECT_EQ(FeatureStore::getFeatureDataLen(0xC5), 4096);

    EXPECT_THROW(
        FeatureStore::getFeatureDataLen(0x00),
        sdbusplus::xyz::openbmc_project::Common::Error::InvalidArgument);
    EXPECT_THROW(
        FeatureStore::getFeatureDataLen(0x03),
        sdbusplus::xyz::openbmc_project::Common::Error::InvalidArgument);
    EXPECT_THROW(
        FeatureStore::getFeatureDataLen(0xFF),
        sdbusplus::xyz::openbmc_project::Common::Error::InvalidArgument);
}

int main(int argc, char** argv)
{
    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}