#include "NVMeMetricStore.hpp"

#include <fcntl.h>
#include <unistd.h>

#include <boost/asio.hpp>
#include <boost/endian/arithmetic.hpp>

#include <cstring>

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

std::vector<std::shared_ptr<MetricBase<>>>
    MetricStore::getVendorMatrics(const std::string& /* path */)
{
    return {};
}

class MetricMock : public MetricBase<std::chrono::steady_clock>
{
  public:
    MetricMock(std::shared_ptr<Scheduler<std::chrono::steady_clock>> scheduler,
               std::string_view id) :
        MetricBase<std::chrono::steady_clock>(
            std::move(scheduler), std::chrono::steady_clock::duration::max()),
        id(id)
    {}

    MOCK_METHOD(
        void, readDevice,
        (std::function<void(std::error_code ec, size_t size, bool complete)> &&
         cb),
        (noexcept, override));

    std::string_view getIdentifier() const noexcept override
    {
        return id;
    }

    MOCK_METHOD(bool, isCacheValid, (), (const, noexcept, override));

    MOCK_METHOD((std::tuple<std::chrono::time_point<ClockType>,
                            std::chrono::time_point<ClockType>,
                            std::span<const uint8_t>>),
                getCache, (), (const, noexcept, override));

  private:
    std::string id;
};

class TestNVMeRefresh : public testing::Test
{
  public:
    boost::asio::io_context io;
    std::shared_ptr<Scheduler<std::chrono::steady_clock>> scheduler;
    std::shared_ptr<sdbusplus::asio::connection> conn;

    TestNVMeRefresh()
    {
        scheduler = std::make_shared<Scheduler<std::chrono::steady_clock>>(io);
        conn = std::make_shared<sdbusplus::asio::connection>(io);
        scheduler->start();
        scheduler->dequeue();
    }

    static sdbusplus::message::unix_fd refreshMetric(MetricStore& store,
                                                     std::string name)
    {
        return store.refreshMetric(std::move(name));
    }
};

TEST_F(TestNVMeRefresh, RefreshEmptyName)
{
    MetricStore store(io, conn, "/test/path");
    EXPECT_THROW(
        refreshMetric(store, ""),
        sdbusplus::xyz::openbmc_project::Common::Error::InvalidArgument);
}

TEST_F(TestNVMeRefresh, RefreshUnknownName)
{
    MetricStore store(io, conn, "/test/path");
    EXPECT_THROW(
        refreshMetric(store, "UnknownMetric"),
        sdbusplus::xyz::openbmc_project::Common::Error::InvalidArgument);
}

TEST_F(TestNVMeRefresh, RefreshSuccess)
{
    auto mtx = std::make_shared<MetricMock>(scheduler, "TestMetric");
    testing::Mock::AllowLeak(mtx.get());
    std::vector<std::shared_ptr<MetricBase<std::chrono::steady_clock>>>
        metrics = {mtx};

    {
        MetricStore store(io, conn, "/test/path", metrics.begin(),
                          metrics.end());

        std::chrono::steady_clock::time_point requestTime =
            std::chrono::steady_clock::now();
        // Simulation time: startTime is 5 seconds after requestTime, finishTime
        // is 10 seconds after requestTime
        std::chrono::steady_clock::time_point startTime =
            requestTime + std::chrono::seconds(5);
        std::chrono::steady_clock::time_point finishTime =
            requestTime + std::chrono::seconds(10);
        std::vector<uint8_t> data = {1, 2, 3};

        EXPECT_CALL(*mtx, getCache())
            .WillRepeatedly(testing::Return(std::make_tuple(
                startTime, finishTime,
                std::span<const uint8_t>(data.begin(), data.end()))));

        EXPECT_CALL(*mtx, readDevice(testing::_))
            .WillOnce(testing::Invoke(
                [](std::function<void(std::error_code ec, size_t size,
                                      bool complete)>&& cb) {
            cb({}, 3, true);
        }));

        sdbusplus::message::unix_fd fd = refreshMetric(store, "TestMetric");
        int readFd = fd.fd;
        EXPECT_GE(readFd, 0);

        // Run IO to trigger the refresh and subsequent write to pipe
        bool dataRead = false;
        std::vector<uint8_t> readBuffer(sizeof(MetricHeader) + data.size());

        auto asyncPipe =
            std::make_shared<boost::asio::posix::stream_descriptor>(io, readFd);
        boost::asio::async_read(
            *asyncPipe, boost::asio::buffer(readBuffer),
            [&](const boost::system::error_code& ec, std::size_t n) {
            EXPECT_FALSE(ec);
            EXPECT_EQ(n, readBuffer.size());
            dataRead = true;
        });

        while (!dataRead)
        {
            io.poll();
        }

        MetricHeader header{};
        std::memcpy(&header, readBuffer.data(), sizeof(header));
        EXPECT_EQ(header.lens, boost::endian::little_uint32_t(sizeof(header)));
        EXPECT_EQ(header.startTime, boost::endian::little_uint64_t(
                                        startTime.time_since_epoch().count()));
        EXPECT_EQ(header.finishTime,
                  boost::endian::little_uint64_t(
                      finishTime.time_since_epoch().count()));

        std::vector<uint8_t> actualData(readBuffer.begin() + sizeof(header),
                                        readBuffer.end());
        EXPECT_EQ(actualData, data);
    }

    // Explicitly clear remaining refs
    mtx.reset();
    metrics.clear();
    scheduler.reset();
}

TEST_F(TestNVMeRefresh, RefreshDouble)
{
    auto mtx = std::make_shared<MetricMock>(scheduler, "TestMetric");
    testing::Mock::AllowLeak(mtx.get());
    std::vector<std::shared_ptr<MetricBase<std::chrono::steady_clock>>>
        metrics = {mtx};

    MetricStore store(io, conn, "/test/path", metrics.begin(), metrics.end());

    // 1. Simulate an ongoing refresh that started BEFORE the request
    std::chrono::steady_clock::time_point preRequestTime =
        std::chrono::steady_clock::now() - std::chrono::seconds(10);
    std::vector<uint8_t> data = {1, 2, 3};

    EXPECT_CALL(*mtx, readDevice(testing::_))
        .Times(2)
        .WillRepeatedly(testing::Invoke(
            [](std::function<void(std::error_code ec, size_t size,
                                  bool complete)>&& cb) { cb({}, 3, true); }));

    std::chrono::steady_clock::time_point secondStartTime;
    std::chrono::steady_clock::time_point secondFinishTime;

    EXPECT_CALL(*mtx, getCache())
        .WillOnce(testing::Return(std::make_tuple(
            preRequestTime, preRequestTime,
            std::span<const uint8_t>(data.begin(), data.end()))))
        .WillRepeatedly(testing::Invoke([&]() {
        secondStartTime = std::chrono::steady_clock::now() +
                          std::chrono::seconds(5);
        secondFinishTime = secondStartTime + std::chrono::seconds(5);
        return std::make_tuple(
            secondStartTime, secondFinishTime,
            std::span<const uint8_t>(data.begin(), data.end()));
    }));

    sdbusplus::message::unix_fd fd = refreshMetric(store, "TestMetric");
    int readFd = fd.fd;
    EXPECT_GE(readFd, 0);

    bool dataRead = false;
    std::vector<uint8_t> readBuffer(sizeof(MetricHeader) + data.size());
    auto asyncPipe =
        std::make_shared<boost::asio::posix::stream_descriptor>(io, readFd);
    boost::asio::async_read(
        *asyncPipe, boost::asio::buffer(readBuffer),
        [&](const boost::system::error_code& ec, std::size_t n) {
        EXPECT_FALSE(ec);
        EXPECT_EQ(n, readBuffer.size());
        dataRead = true;
    });

    while (!dataRead)
    {
        io.poll();
    }

    MetricHeader header{};
    std::memcpy(&header, readBuffer.data(), sizeof(header));
    EXPECT_EQ(header.startTime,
              boost::endian::little_uint64_t(
                  secondStartTime.time_since_epoch().count()));
}

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