#include "NVMeDevice.hpp"
#include "NVMeIntf.hpp"
#include "NVMeMi.hpp"
#include "Utils.hpp"

#include <boost/asio/steady_timer.hpp>

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

class MockMctpEndpoint : public MctpEndpoint
{
  public:
    MOCK_METHOD(int, network, (), (const, override));
    MOCK_METHOD(uint8_t, eid, (), (const, override));
    MOCK_METHOD(void, subscribe,
                (MctpEndpoint::Event && degraded,
                 MctpEndpoint::Event&& available,
                 MctpEndpoint::Event&& removed),
                (override));
    MOCK_METHOD(void, setMtu,
                (uint32_t mtu,
                 std::function<void(const std::error_code& ec)>&& completed),
                (override));
    MOCK_METHOD(void, remove, (), (override));
    MOCK_METHOD(void, recover, (), (override));
    MOCK_METHOD(std::string, describe, (), (const, override));
    MOCK_METHOD(bool, isI2cAccessible, (), (const, override));
};

class MockMctpDevice : public MctpDevice
{
  public:
    MOCK_METHOD(void, setup,
                (std::function<void(const std::error_code& ec,
                                    const std::shared_ptr<MctpEndpoint>& ep)> &&
                 action),
                (override));
    MOCK_METHOD(void, remove, (), (override));
    MOCK_METHOD(std::string, describe, (), (const, override));
};

class TestNVMeMi : public NVMeMi
{
  public:
    using NVMeMi::NVMeMi;
    bool isStarted = false;

    void start(const std::shared_ptr<MctpEndpoint>& ep) override
    {
        isStarted = true;
        NVMeMi::start(ep);
    }
    void stop() override
    {
        isStarted = false;
        NVMeMi::stop();
    }
};

TEST(NVMeRecovery, optimisationFailure)
{
    boost::asio::io_context io;

    auto mctpEp = std::make_shared<MockMctpEndpoint>();
    EXPECT_CALL(*mctpEp, describe())
        .WillRepeatedly(testing::Return("Mock MCTP Endpoint"));
    EXPECT_CALL(*mctpEp, eid()).WillRepeatedly(testing::Return(9));
    EXPECT_CALL(*mctpEp, network()).WillRepeatedly(testing::Return(1));
    EXPECT_CALL(*mctpEp, isI2cAccessible())
        .WillRepeatedly(testing::Return(true));

    boost::asio::spawn(io, [&](boost::asio::yield_context yield) {
        MctpEndpoint::Event degradedHandler;
        MctpEndpoint::Event availableHandler;
        MctpEndpoint::Event removedHandler;
        EXPECT_CALL(*mctpEp, subscribe(testing::_, testing::_, testing::_))
            .WillOnce(testing::DoAll(testing::SaveArg<0>(&degradedHandler),
                                     testing::SaveArg<1>(&availableHandler),
                                     testing::SaveArg<2>(&removedHandler),
                                     testing::InvokeWithoutArgs([&]() {
            io.post([&]() { availableHandler(mctpEp); });
        })));
        EXPECT_CALL(*mctpEp, recover())
            .Times(testing::Between(1, 3))
            .WillRepeatedly(testing::InvokeWithoutArgs([&]() {
            io.post([&]() { degradedHandler(mctpEp); });
            io.post([&]() { availableHandler(mctpEp); });
        }));

        boost::asio::steady_timer timer(io);
        auto mctpDev = std::make_shared<MockMctpDevice>();
        EXPECT_CALL(*mctpDev, describe())
            .WillRepeatedly(testing::Return("Mock MCTP Device"));
        EXPECT_CALL(*mctpDev, setup(testing::_))
            .WillOnce(testing::InvokeArgument<0>(std::error_code(), mctpEp));

        auto systemBus = std::make_shared<sdbusplus::asio::connection>(io);
        auto worker = NVMeMiWorker::create(io);
        auto intf = NVMeIntf::create<NVMeMi>(io, systemBus, mctpDev, worker);
        auto nvmeDev = NVMeDevice::create(io, "bar", mctpDev, intf,
                                          std::chrono::seconds(2));
        nvmeDev->start();
        timer.expires_after(std::chrono::seconds(6));
        timer.async_wait(yield);
        // io.run_for(std::chrono::seconds(6));

        nvmeDev->stop();
        // io.run_for(std::chrono::seconds(1));
        timer.expires_after(std::chrono::seconds(1));
        timer.async_wait(yield);

        // https://stackoverflow.com/a/10289205
        EXPECT_TRUE(testing::Mock::VerifyAndClearExpectations(mctpEp.get()));
        EXPECT_TRUE(testing::Mock::VerifyAndClearExpectations(mctpDev.get()));
        io.stop();
    });

    io.run();
}

TEST(NVMeRecovery, tightEventCycle)
{
    boost::asio::io_context io;

    auto mctpEp = std::make_shared<MockMctpEndpoint>();
    EXPECT_CALL(*mctpEp, describe())
        .WillRepeatedly(testing::Return("Mock MCTP Endpoint"));
    EXPECT_CALL(*mctpEp, eid()).WillRepeatedly(testing::Return(9));
    EXPECT_CALL(*mctpEp, network()).WillRepeatedly(testing::Return(1));
    EXPECT_CALL(*mctpEp, isI2cAccessible())
        .WillRepeatedly(testing::Return(true));

    boost::asio::spawn(io, [&](boost::asio::yield_context yield) {
        MctpEndpoint::Event degradedHandler;
        MctpEndpoint::Event availableHandler;
        MctpEndpoint::Event removedHandler;
        EXPECT_CALL(*mctpEp, subscribe(testing::_, testing::_, testing::_))
            .WillOnce(testing::DoAll(testing::SaveArg<0>(&degradedHandler),
                                     testing::SaveArg<1>(&availableHandler),
                                     testing::SaveArg<2>(&removedHandler),
                                     testing::InvokeWithoutArgs([&]() {
            io.post([&]() { availableHandler(mctpEp); });
        })));
        EXPECT_CALL(*mctpEp, recover()).Times(testing::AnyNumber());

        boost::asio::steady_timer timer(io);
        auto mctpDev = std::make_shared<MockMctpDevice>();
        EXPECT_CALL(*mctpDev, describe())
            .WillRepeatedly(testing::Return("Mock MCTP Device"));
        EXPECT_CALL(*mctpDev, setup(testing::_))
            .WillOnce(testing::InvokeArgument<0>(std::error_code(), mctpEp));

        auto systemBus = std::make_shared<sdbusplus::asio::connection>(io);
        auto worker = NVMeMiWorker::create(io);
        auto intf = NVMeIntf::create<TestNVMeMi>(io, systemBus, mctpDev,
                                                 worker);
        auto testMi = std::dynamic_pointer_cast<TestNVMeMi>(
            std::get<std::shared_ptr<NVMeMiIntf>>(intf.getInferface()));
        auto nvmeDev = NVMeDevice::create(io, "bar", mctpDev, intf,
                                          std::chrono::seconds(2));
        nvmeDev->start();

        // Wait for initial available handler to fire
        timer.expires_after(std::chrono::milliseconds(100));
        timer.async_wait(yield);

        // 1. Degrade
        degradedHandler(mctpEp);

        // 2. Wait < 1 sec
        timer.expires_after(std::chrono::milliseconds(500));
        timer.async_wait(yield);

        // 3. Available
        availableHandler(mctpEp);

        // 4. Wait < 1 sec again (before 2s grace period ends)
        timer.expires_after(std::chrono::milliseconds(500));
        timer.async_wait(yield);

        // 5. Degrade
        degradedHandler(mctpEp);

        // Let things settle to ensure no crashes
        timer.expires_after(std::chrono::seconds(2));
        timer.async_wait(yield);

        // Validate the final state for the mctpEp should be degraded
        EXPECT_FALSE(testMi->isStarted);

        nvmeDev->stop();

        timer.expires_after(std::chrono::seconds(1));
        timer.async_wait(yield);

        EXPECT_TRUE(testing::Mock::VerifyAndClearExpectations(mctpEp.get()));
        EXPECT_TRUE(testing::Mock::VerifyAndClearExpectations(mctpDev.get()));
        io.stop();
    });

    io.run();
}

TEST(NVMeRecovery, consecutiveAvailableBypass)
{
    boost::asio::io_context io;

    auto mctpEp = std::make_shared<MockMctpEndpoint>();
    EXPECT_CALL(*mctpEp, describe())
        .WillRepeatedly(testing::Return("Mock MCTP Endpoint"));
    EXPECT_CALL(*mctpEp, eid()).WillRepeatedly(testing::Return(9));
    EXPECT_CALL(*mctpEp, network()).WillRepeatedly(testing::Return(1));
    EXPECT_CALL(*mctpEp, isI2cAccessible())
        .WillRepeatedly(testing::Return(true));

    boost::asio::spawn(io, [&](boost::asio::yield_context yield) {
        MctpEndpoint::Event degradedHandler;
        MctpEndpoint::Event availableHandler;
        MctpEndpoint::Event removedHandler;
        EXPECT_CALL(*mctpEp, subscribe(testing::_, testing::_, testing::_))
            .WillOnce(testing::DoAll(testing::SaveArg<0>(&degradedHandler),
                                     testing::SaveArg<1>(&availableHandler),
                                     testing::SaveArg<2>(&removedHandler),
                                     testing::InvokeWithoutArgs([&]() {
            io.post([&]() { availableHandler(mctpEp); });
        })));
        EXPECT_CALL(*mctpEp, recover()).Times(testing::AnyNumber());

        boost::asio::steady_timer timer(io);
        auto mctpDev = std::make_shared<MockMctpDevice>();
        EXPECT_CALL(*mctpDev, describe())
            .WillRepeatedly(testing::Return("Mock MCTP Device"));
        EXPECT_CALL(*mctpDev, setup(testing::_))
            .WillOnce(testing::InvokeArgument<0>(std::error_code(), mctpEp));

        auto systemBus = std::make_shared<sdbusplus::asio::connection>(io);
        auto worker = NVMeMiWorker::create(io);
        auto intf = NVMeIntf::create<TestNVMeMi>(io, systemBus, mctpDev,
                                                 worker);
        auto testMi = std::dynamic_pointer_cast<TestNVMeMi>(
            std::get<std::shared_ptr<NVMeMiIntf>>(intf.getInferface()));
        auto nvmeDev = NVMeDevice::create(io, "bar", mctpDev, intf,
                                          std::chrono::seconds(2));
        nvmeDev->start();

        // Wait for initial available handler to fire
        timer.expires_after(std::chrono::milliseconds(100));
        timer.async_wait(yield);

        // 1. Degrade
        degradedHandler(mctpEp);
        EXPECT_FALSE(testMi->isStarted);

        // 2. Available (starts 2s timer)
        availableHandler(mctpEp);
        EXPECT_FALSE(testMi->isStarted);

        // 3. Wait 0.5s, well within the 2s grace period
        timer.expires_after(std::chrono::milliseconds(500));
        timer.async_wait(yield);

        // 4. Consecutive Available
        // The current buggy code bypasses the grace period here and calls
        // start() immediately.
        availableHandler(mctpEp);
        EXPECT_FALSE(testMi->isStarted);

        // 5. Wait for the remainder of the 2s grace period
        timer.expires_after(std::chrono::seconds(2));
        timer.async_wait(yield);

        EXPECT_TRUE(testMi->isStarted);

        nvmeDev->stop();

        timer.expires_after(std::chrono::seconds(1));
        timer.async_wait(yield);

        EXPECT_TRUE(testing::Mock::VerifyAndClearExpectations(mctpEp.get()));
        EXPECT_TRUE(testing::Mock::VerifyAndClearExpectations(mctpDev.get()));
        io.stop();
    });

    io.run();
}

// Unused, but required to link successfully
std::unordered_map<std::string, void*> pluginLibMap = {};

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