#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),
        (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));
};

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

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

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

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);

    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, [&](boost::asio::yield_context yield) {
        result = FeatureStoreTest::callGetFeature(
            featureStore, std::move(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);

    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, [&](boost::asio::yield_context yield) {
        FeatureStoreTest::callSetFeature(featureStore, std::move(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
        }
    });

    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, 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, [&](boost::asio::yield_context yield) {
        try
        {
            FeatureStoreTest::callGetFeature(featureStore, std::move(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);

    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, [&](boost::asio::yield_context yield) {
        try
        {
            FeatureStoreTest::callGetFeature(featureStore, std::move(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, [&](boost::asio::yield_context yield) {
        try
        {
            FeatureStoreTest::callSetFeature(featureStore, std::move(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);

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

    bool threw = false;
    bool done = false;
    boost::asio::spawn(io, [&](boost::asio::yield_context yield) {
        try
        {
            FeatureStoreTest::callSetFeature(featureStore, std::move(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);

    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, [&](boost::asio::yield_context yield) {
        try
        {
            FeatureStoreTest::callSetFeature(featureStore, std::move(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);

    // 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, [&](boost::asio::yield_context yield) {
        try
        {
            FeatureStoreTest::callSetFeature(featureStore, std::move(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, [&](boost::asio::yield_context yield) {
        try
        {
            FeatureStoreTest::callSetFeature(featureStore, std::move(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);
}

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