#include "cper_encoder.hpp"
#include "mock_cper_encoder.hpp"

#include <chrono>
#include <concepts>

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

using testing::_;
using testing::Contains;

namespace uefi::cper
{
template <typename T>
concept StringOrVector =
    std::is_same_v<T, std::string> || std::is_same_v<T, std::vector<uint8_t>>;

class CperEncoderTest : public testing::Test
{
  public:
    void SetUp() override
    {}

    bool isArrayEq(const void* arr1, const void* arr2, const uint64_t arr1Size,
                   const uint64_t arr2Size)
    {
        if (arr1Size != arr2Size)
        {
            return false;
        }
        return memcmp(arr1, arr2, arr1Size) == 0;
    }

    template <StringOrVector CperType>
    void expectValidHeader(const CperType& cperArray,
                           const uint16_t numSections, const uint64_t cperSize,
                           const uint64_t timestamp)
    {
        RecordHeader header;
        memcpy(&header, cperArray.data(), sizeof(header));

        EXPECT_EQ(header.signature, kRecordSignature);
        EXPECT_EQ(header.revision, kRecordRevision);
        EXPECT_EQ(header.signatureEnd, kRecordSignatureEnd);
        EXPECT_EQ(header.sectionCount, numSections);
        EXPECT_EQ(header.errorSeverity, kMockRecordSeverity);
        EXPECT_EQ(header.validationBits, kMockValidationBits);
        EXPECT_EQ(header.recordLength, cperSize);
        EXPECT_EQ(header.timestamp, timestamp);
        EXPECT_EQ(header.platformId, kMockPlatformId);
        EXPECT_EQ(header.partitionId, kMockPartitionId);
        EXPECT_EQ(header.creatorId, kMockCreatorId);
        EXPECT_EQ(header.notificationType, kMockNotificationType);
        EXPECT_EQ(header.recordId, kMockRecordId);
        EXPECT_EQ(header.flags, kMockHeaderFlags);
        EXPECT_EQ(header.persistenceInformation, kMockPersistenceInformation);
        EXPECT_THAT(header.reserved, Contains(0).Times(header.reserved.size()));
    }

    template <StringOrVector CperType>
    void expectValidSectionDescriptor(const CperType& cperArray,
                                      const uint64_t sectionNum,
                                      const uint32_t expectedOffset,
                                      const uint32_t expectedLength,
                                      const uint32_t expectedFlag,
                                      const guid_t expectedSectionType)
    {
        const uint64_t sectionOffset =
            kRecordHeaderSizeBytes + (kSectionDescriptorSizeBytes * sectionNum);

        SectionDescriptor descriptor;
        memcpy(&descriptor, &cperArray[sectionOffset], sizeof(descriptor));

        EXPECT_EQ(descriptor.sectionOffset, expectedOffset);
        EXPECT_EQ(descriptor.sectionLength, expectedLength);
        EXPECT_EQ(descriptor.revision, kMockSectionRevision);
        EXPECT_EQ(descriptor.validationBits, kMockSectionValidationBits);
        EXPECT_EQ(descriptor.reserved, 0);
        EXPECT_EQ(descriptor.flags, expectedFlag);
        EXPECT_EQ(descriptor.sectionType, expectedSectionType);
        EXPECT_EQ(descriptor.fruId, kMockFruId);
        EXPECT_EQ(descriptor.sectionSeverity, kMockSectionSeverity);
        EXPECT_EQ(descriptor.fruText, kMockFruText);
    }

    template <ByteLike BodyType, StringOrVector CperType>
    void expectValidSectionBody(
        const CperType& cperArray,
        const std::span<const BodyType> expectedSectionBody,
        const uint64_t totalSections, const uint32_t sectionOffset)
    {
        // The offset set in the section descriptors is the offset in bytes
        // after the header and descriptors.
        const uint32_t kTrueBodyOffset =
            sectionOffset + kRecordHeaderSizeBytes +
            (kSectionDescriptorSizeBytes * totalSections);

        const uint64_t kSectionBodySize = expectedSectionBody.size();

        ASSERT_GE(cperArray.size(), (kTrueBodyOffset + kSectionBodySize));

        CperType sectionBody(kSectionBodySize, 0);
        memcpy(sectionBody.data(), &cperArray[kTrueBodyOffset],
               sectionBody.size());

        EXPECT_TRUE(isArrayEq(sectionBody.data(), expectedSectionBody.data(),
                              sectionBody.size(), expectedSectionBody.size()));
    }

    void expectEqTimestamp(
        const uint64_t cperTimestamp,
        const std::chrono::time_point<std::chrono::system_clock> ref,
        const bool isPrecise = true)
    {
        Timestamp structTimestamp;
        memcpy(&structTimestamp, &cperTimestamp, sizeof(structTimestamp));

        time_t tt = std::chrono::system_clock::to_time_t(ref);
        tm utcTime = *gmtime(&tt);

        EXPECT_EQ(structTimestamp.seconds, utcTime.tm_sec);
        EXPECT_EQ(structTimestamp.minutes, utcTime.tm_min);
        EXPECT_EQ(structTimestamp.hours, utcTime.tm_hour);
        EXPECT_EQ(structTimestamp.isPrecise, (isPrecise & 0x1));
        EXPECT_EQ(structTimestamp.day, utcTime.tm_mday);
        EXPECT_EQ(structTimestamp.month, utcTime.tm_mon);
        EXPECT_EQ(structTimestamp.year, (utcTime.tm_year + 1900) % 100);
        EXPECT_EQ(structTimestamp.century, (utcTime.tm_year + 1900) / 100);
    }
};

TEST_F(CperEncoderTest, TooMuchData)
{
    MockCperEncoder<uint8_t> mockEncoder;
    uint64_t totalSize = kRecordHeaderSizeBytes;

    // Exceed the 4GiB CPER Log max size by adding 2^16 * 2^16 Bytes worth of
    // section bodies along with their descriptors.
    for (uint64_t i = 0; i < std::numeric_limits<uint16_t>::max(); i++)
    {
        std::array<uint8_t, std::numeric_limits<uint16_t>::max()> sectionBody;
        totalSize += sectionBody.size() + kSectionDescriptorSizeBytes;

        // Note the encoder does not own the sectionBody, it just tracks the
        // body ptrs and size. This would never work in a real application since
        // the sectionBody only exists within each loops iteration.
        if (totalSize < std::numeric_limits<uint32_t>::max())
        {

            EXPECT_TRUE(mockEncoder.addSection(sectionBody, kMockSectionFlag,
                                               kMockSectionType));
        }
        else
        {
            EXPECT_FALSE(mockEncoder.addSection(sectionBody, kMockSectionFlag,
                                                kMockSectionType));
            return;
        }
    }

    FAIL() << "Should never be able to complete the loop";
}

TEST_F(CperEncoderTest, EncodeMultipleCharSectionsToArrayAndString)
{
    MockCperEncoder<char> mockEncoder;

    // Fill our CPER Log with some quotes to get through the work week.
    std::vector<std::string> sectionBodies;
    sectionBodies.push_back("The first principle is that you must not fool "
                            "yourself and you are the easiest person to fool.");
    sectionBodies.push_back("Spread love everywhere you go.");
    sectionBodies.push_back("The only thing we have to fear is fear itself.");
    sectionBodies.push_back("If you can't explain it to a six year old, you "
                            "don't understand it yourself.");
    sectionBodies.push_back("I have no wings, so I guess I'll look up at this "
                            "sky, and crawl along the Earth.");
    sectionBodies.push_back("If I have seen further, it is by standing on the "
                            "shoulders of giants.");

    const uint16_t kNumSections = sectionBodies.size();

    // Add the sections to our encoder.
    uint64_t totalSectionBodySize = 0;
    for (const std::string& body : sectionBodies)
    {
        mockEncoder.addSection(body, kMockSectionFlag, kMockSectionType);
        totalSectionBodySize += body.size();
    }

    const uint64_t kExpectedCperSize =
        (kRecordHeaderSizeBytes + (kSectionDescriptorSizeBytes * kNumSections) +
         totalSectionBodySize);

    // Serialize to an array.
    EXPECT_CALL(mockEncoder, mockCreateRecordHeader).Times(1);
    EXPECT_CALL(mockEncoder, mockCreateSectionDescriptor).Times(kNumSections);
    const std::vector<uint8_t> cperArray = mockEncoder.serializeToByteArray();
    ASSERT_EQ(cperArray.size(), kExpectedCperSize);

    // Check that the serialized arrays header contains the same contents we
    // configured the encoder to use.
    expectValidHeader(cperArray, kNumSections, kExpectedCperSize,
                      mockEncoder.timestamp);

    // Check that the section descriptors and bodies all contain the correct
    // information about their respective section.
    uint32_t expectedOffset = 0;
    uint64_t sectionIdx = 0;
    for (const std::string& body : sectionBodies)
    {
        const uint32_t kBodyLength = body.size();
        expectValidSectionDescriptor(cperArray, sectionIdx, expectedOffset,
                                     kBodyLength, kMockSectionFlag,
                                     kMockSectionType);

        expectValidSectionBody<char>(cperArray, body, kNumSections,
                                     expectedOffset);

        expectedOffset += kBodyLength;
        sectionIdx++;
    }

    // Serialize to a string.
    EXPECT_CALL(mockEncoder, mockCreateRecordHeader).Times(1);
    EXPECT_CALL(mockEncoder, mockCreateSectionDescriptor).Times(kNumSections);
    const std::string cperStr = mockEncoder.serializeToString();
    ASSERT_EQ(cperStr.size(), kExpectedCperSize);

    // Check that the serialized arrays header contains the same contents we
    // configured the encoder to use.
    expectValidHeader(cperStr, kNumSections, kExpectedCperSize,
                      mockEncoder.timestamp);

    // Check that the section descriptors and bodies all contain the correct
    // information about their respective section.
    expectedOffset = 0;
    sectionIdx = 0;
    for (const std::string& body : sectionBodies)
    {
        const uint32_t kBodyLength = body.size();
        expectValidSectionDescriptor(cperStr, sectionIdx, expectedOffset,
                                     kBodyLength, kMockSectionFlag,
                                     kMockSectionType);

        expectValidSectionBody<char>(cperStr, body, kNumSections,
                                     expectedOffset);

        expectedOffset += kBodyLength;
        sectionIdx++;
    }
}

TEST_F(CperEncoderTest, EncodeOneU8SectionToString)
{
    MockCperEncoder<uint8_t> mockEncoder;

    const std::array<uint8_t, 10> sectionBody = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
    const uint64_t kSectionBodySize = sectionBody.size();
    constexpr uint16_t kNumSections = 1;

    // Add 1 section
    mockEncoder.addSection(sectionBody, kMockSectionFlag, kMockSectionType);

    // Serialize the data into a CPER byte array.
    EXPECT_CALL(mockEncoder, mockCreateRecordHeader).Times(1);
    EXPECT_CALL(mockEncoder, mockCreateSectionDescriptor).Times(kNumSections);
    const std::string cperArray = mockEncoder.serializeToString();

    // If this bricks, abort ship - there is likely memory corruption.
    const uint64_t kExpectedCperSize =
        (kSectionBodySize + kSectionDescriptorSizeBytes +
         kRecordHeaderSizeBytes);

    ASSERT_EQ(cperArray.size(), kExpectedCperSize);

    // Check that the serialized arrays header contains the same contents we
    // configured the encoder to use.
    expectValidHeader(cperArray, kNumSections, kExpectedCperSize,
                      mockEncoder.timestamp);

    // Check that the section descriptor contains the correct information about
    // the section.
    const uint32_t kExpectedOffset = 0;
    expectValidSectionDescriptor(cperArray, 0, kExpectedOffset,
                                 kSectionBodySize, kMockSectionFlag,
                                 kMockSectionType);

    expectValidSectionBody<uint8_t>(cperArray, sectionBody, kNumSections,
                                    kExpectedOffset);
}

TEST_F(CperEncoderTest, EncodeOneCharSectionToArray)
{
    MockCperEncoder<char> mockEncoder;

    const std::string sectionBody =
        "Hello world, this is a simplistic example of a CPER Section Body.";
    const uint64_t kSectionBodySize = sectionBody.size();
    constexpr uint16_t kNumSections = 1;

    // Add 1 section
    mockEncoder.addSection(sectionBody, kMockSectionFlag, kMockSectionType);

    // Serialize the data into a CPER byte array.
    EXPECT_CALL(mockEncoder, mockCreateRecordHeader).Times(1);
    EXPECT_CALL(mockEncoder, mockCreateSectionDescriptor).Times(kNumSections);
    const std::vector<uint8_t> cperArray = mockEncoder.serializeToByteArray();

    // If this bricks, abort ship - there is likely memory corruption.
    const uint64_t kExpectedCperSize =
        (kSectionBodySize + kSectionDescriptorSizeBytes +
         kRecordHeaderSizeBytes);

    ASSERT_EQ(cperArray.size(), kExpectedCperSize);

    // Check that the serialized arrays header contains the same contents we
    // configured the encoder to use.
    expectValidHeader(cperArray, kNumSections, kExpectedCperSize,
                      mockEncoder.timestamp);

    // Check that the section descriptor contains the correct information about
    // the section.
    const uint32_t kExpectedOffset = 0;
    expectValidSectionDescriptor(cperArray, 0, kExpectedOffset,
                                 kSectionBodySize, kMockSectionFlag,
                                 kMockSectionType);

    expectValidSectionBody<char>(cperArray, sectionBody, kNumSections,
                                 kExpectedOffset);
}

TEST_F(CperEncoderTest, SeqMemcopy)
{
    std::string section1 = "Hello";
    std::string section2 = " ";
    std::string section3 = "World";

    std::string dest(section1.size() + section2.size() + section3.size(), 0);

    SeqMemcopy<char> seqMemcpy(dest);

    EXPECT_NO_THROW(seqMemcpy.copy(section1.data(), section1.size()));
    EXPECT_NO_THROW(seqMemcpy.copy(section2.data(), section2.size()));
    EXPECT_NO_THROW(seqMemcpy.copy(section3.data(), section3.size()));

    EXPECT_EQ(dest, section1 + section2 + section3);

    std::string section4 = "This is an attempted buffer overflow!";

    EXPECT_THROW(seqMemcpy.copy(section4.data(), section4.size()),
                 std::out_of_range);
}

TEST_F(CperEncoderTest, CreateCperTimestamp)
{
    const auto now = std::chrono::system_clock::now();

    const uint64_t preciseTimestamp = createCperTimestamp(now);
    expectEqTimestamp(preciseTimestamp, now);

    const uint64_t impreciseTimestamp = createCperTimestamp(now, false);
    expectEqTimestamp(impreciseTimestamp, now, false);
}

TEST_F(CperEncoderTest, SmokeTest)
{
    RecordHeader recordHeader;

    const std::array<uint8_t, 4> kExpectedSignature = {'C', 'P', 'E', 'R'};
    const std::array<uint8_t, 12> kExpectedReserved = {0};

    EXPECT_EQ(recordHeader.signature, kExpectedSignature);
    EXPECT_EQ(recordHeader.revision, 0x0000);
    EXPECT_EQ(recordHeader.signatureEnd, 0xffffffff);
    EXPECT_EQ(recordHeader.reserved, kExpectedReserved);

    // Check for proper packing by putting a needle in the haystack.
    recordHeader.flags = 0xBEEF;
    const size_t flagsOffset = 104;

    std::array<uint8_t, sizeof(recordHeader)> headerBin;
    memcpy(&headerBin, &recordHeader, sizeof(headerBin));

    uint32_t flagsValue;
    memcpy(&flagsValue, &headerBin[flagsOffset], sizeof(flagsValue));
    EXPECT_EQ(recordHeader.flags, flagsValue);
}
} // namespace uefi::cper