// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "mtd_util_unittest.hpp"

#include "fs.hpp"
#include "mtd_util.hpp"

#include <fcntl.h>
#include <mtd/mtd-user.h>

#include <cstring>
#include <exception>
#include <filesystem>
#include <map>
#include <optional>
#include <random>
#include <vector>

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

// NOLINTNEXTLINE(google-build-using-namespace)
using namespace std::literals;
namespace fs = std::filesystem;

using ::testing::_;
using ::testing::AllOf;
using ::testing::ContainerEq;
using ::testing::DoAll;
using ::testing::Field;
using ::testing::Invoke;
using ::testing::MatcherCast;
using ::testing::NotNull;
using ::testing::Return;
using ::testing::SetErrnoAndReturn;
using ::testing::StrEq;

auto static const test_str = "Hello, world!"s;
std::vector<uint8_t> static const test_buf(test_str.begin(), test_str.end());

auto constexpr testName = "test-name";

std::map<fs::path, std::optional<std::string>> mtdMap;

std::vector<fs::path>
google::hoth::internal::listDir(const fs::path& /*unused*/)
{
    std::vector<fs::path> output;

    output.reserve(mtdMap.size());
    for (const auto& p : mtdMap) {
      output.push_back(p.first);
    }

    return output;
}

std::string google::hoth::internal::getFirstLine(const fs::path& path)
{
    if (path.filename() != "name")
    {
        throw std::ios_base::failure("file not found");
    }

    auto result = mtdMap.at(path.parent_path().string());

    try
    {
        return result.value();
    }
    catch (const std::bad_optional_access& e)
    {
        throw std::ios_base::failure("file not found");
    }
}

// Custom action to typecast from void* to mtd_info_t* in our mock return
ACTION_P(SetArg2ToMtdPointer, value)
{
    *static_cast<mtd_info_t*>(arg2) = value;
}

TEST(PartitionTest, emptyMtdMapThrowsException)
{
    // Attempt to findPartition when mtdMap is empty, expect exception thrown
    EXPECT_THROW(google::hoth::internal::mtdImpl.findPartition(testName),
                 std::exception);
}

TEST(PartitionTest, mtdMapDoesNotContainPathToFileThrowsException)
{
    // Attempt to findPartition when mtdMap does not contain path to testName
    // file, expect exception thrown
    mtdMap = {{"/sys/class/mtd/mtd0", "b"}, {"/sys/class/mtd/mtd1", "c"}};

    EXPECT_THROW(google::hoth::internal::mtdImpl.findPartition(testName),
                 std::exception);
}

TEST(PartitionTest, mtdMapContainsPathToFileReturnsPath)
{
    // Attempt to findPartition when mtdMap contains a path with testName file
    // and another path without, expect the path that contains the testName
    // file to be returned
    mtdMap = {{"/sys/class/mtd/mtd0", "b"}, {"/sys/class/mtd/mtd1", testName}};

    EXPECT_EQ("/dev/mtd1",
              google::hoth::internal::mtdImpl.findPartition(testName));
}

TEST(PartitionTest, mtdMapContainsPathToFileAndPathsWithoutFilesReturnsPath)
{
    // Attempt to findPartition when mtdMap contains a path with testName file,
    // path with a random file and mtd*ro folders which don't have a "name"
    // file and will throw an exception when trying to read them. We expect
    // the path that contains the testName file to be returned.
    mtdMap = {{"/sys/class/mtd/mtd0", "b"},
              {"/sys/class/mtd/mtd0ro", std::nullopt},
              {"/sys/class/mtd/mtd1", testName},
              {"/sys/class/mtd/mtd1ro", std::nullopt}};

    EXPECT_EQ("/dev/mtd1",
              google::hoth::internal::mtdImpl.findPartition(testName));
}

TEST_F(FlashCopyTest, openInvalidPartitionThrowsException)
{
    // Attempt to flashCopy to an invalid partition, mock the return of the
    // sys open call with ENOENT which is "No such file or directory".
    // Expect exception thrown.
    std::vector<uint8_t> data;

    EXPECT_CALL(sys, open(StrEq(TEST_PATH), _))
        .WillOnce(SetErrnoAndReturn(ENOENT, -1));

    EXPECT_THROW(
        google::hoth::internal::mtdImpl.flashCopy(&sys, data, TEST_PATH),
        std::exception);
}

TEST_F(FlashCopyTest, partitionNotMtdThrowsException)
{
    // Attempt to flashCopy to a partition which is not an MTD, mock the return
    // of the ioctl MEMGETINFO call with ENOTTY which is "Inappropriate ioctl
    // for device". Expect exception thrown.
    std::vector<uint8_t> data;

    EXPECT_CALL(sys, open(StrEq(TEST_PATH), _)).WillOnce(Return(TEST_FD));
    EXPECT_CALL(sys, ioctl(TEST_FD, MEMGETINFO, NotNull()))
        .WillOnce(SetErrnoAndReturn(ENOTTY, -1));
    EXPECT_CALL(sys, close(TEST_FD)).WillOnce(Return(0));

    EXPECT_THROW(
        google::hoth::internal::mtdImpl.flashCopy(&sys, data, TEST_PATH),
        std::exception);
}

TEST_F(FlashCopyTest, partitionTooSmallThrowsException)
{
    // Attempt to flashCopy data bigger than the size of partition.
    // Expect exception thrown.
    std::vector<uint8_t> data(10);
    mtd_info_t mtd;
    // Choose mtd size smaller than data size
    mtd.size = 5;

    EXPECT_CALL(sys, open(StrEq(TEST_PATH), _)).WillOnce(Return(TEST_FD));
    EXPECT_CALL(sys, ioctl(TEST_FD, MEMGETINFO, NotNull()))
        .WillOnce(DoAll(SetArg2ToMtdPointer(mtd), Return(0)));
    EXPECT_CALL(sys, close(TEST_FD)).WillOnce(Return(0));

    EXPECT_THROW(
        google::hoth::internal::mtdImpl.flashCopy(&sys, data, TEST_PATH),
        std::exception);
}

TEST_F(FlashCopyTest, eraseNotSupportedThrowsException)
{
    // Attempt to flashCopy with valid mtd and erase_info_t but mock
    // ioctl MEMERASE to return "ENOTSUP" which is "Not Supported", meaning the
    // MTD does not support the memerase functionality. Expect exception thrown.
    std::vector<uint8_t> data(10);
    mtd_info_t mtd;
    // Choose a valid mtd size (value bigger than data size)
    mtd.size = 16;

    EXPECT_CALL(sys, open(StrEq(TEST_PATH), _)).WillOnce(Return(TEST_FD));
    EXPECT_CALL(sys, ioctl(TEST_FD, MEMGETINFO, NotNull()))
        .WillOnce(DoAll(SetArg2ToMtdPointer(mtd), Return(0)));
    EXPECT_CALL(sys, ioctl(TEST_FD, MEMERASE,
                           MatcherCast<void*>(MatcherCast<erase_info_t*>(
                               AllOf(Field(&erase_info_t::start, 0),
                                     Field(&erase_info_t::length, mtd.size))))))
        .WillOnce(SetErrnoAndReturn(ENOTSUP, -1));
    EXPECT_CALL(sys, close(TEST_FD)).WillOnce(Return(0));

    EXPECT_THROW(
        google::hoth::internal::mtdImpl.flashCopy(&sys, data, TEST_PATH),
        std::exception);
}

class FakeMtd
{
  public:
    explicit FakeMtd(uint32_t flags = 0) : flags(flags)
    {
        mtd.type = MTD_NORFLASH;
        mtd.flags = MTD_WRITEABLE;
        mtd.size = deviceSize;
        mtd.erasesize = eraseSize;
        mtd.writesize = 1;

        data.resize(deviceSize);

        if (flags & flagWriteIntr)
        {
            writeIntrRemain = intrCount;
        }
        if (flags & flagReadIntr)
        {
            readIntrRemain = intrCount;
        }
    }

    ssize_t write(int /*unused*/, const void *buf, size_t size)
    {
        // Simulate interruption
        if (writeIntrRemain > 0 && writeOffset != 0)
        {
            writeIntrRemain--;
            errno = EINTR;
            return -1;
        }

        // Simulate I/O error halfway through
        if ((flags & flagWriteFail) && (writeOffset > deviceSize / 2))
        {
            errno = EIO;
            return -1;
        }

        auto writeSize =
            std::min(deviceSize - writeOffset, static_cast<off_t>(size));

        // Simulate short write, but always write at least 1 byte
        if ((flags & flagWriteShort) && writeSize > shortCount)
        {
            writeSize -= shortCount;
        }

        if (writeSize < 0)
        {
            return 0;
        }

        // Only simulate a "write" if flagWriteNoData is false
        if (!(flags & flagWriteNoData))
        {
            std::memcpy(data.data() + writeOffset, buf, writeSize);
        }
        writeOffset += writeSize;

        return writeSize;
    }

    ssize_t read(int /*unused*/, void *buf, size_t size)
    {
        // Simulate interruption
        if (readIntrRemain > 0 && readOffset != 0)
        {
            readIntrRemain--;
            errno = EINTR;
            return -1;
        }

        // Simulate I/O error halfway through
        if ((flags & flagReadFail) && (readOffset > deviceSize / 2))
        {
            errno = EIO;
            return -1;
        }

        auto readSize =
            std::min(deviceSize - readOffset, static_cast<off_t>(size));

        // Simulate short read, but always read at least 1 byte
        if ((flags & flagReadShort) && readSize > shortCount)
        {
            readSize -= shortCount;
        }

        if (readSize < 0)
        {
            return 0;
        }

        // Only simulate a "read" if flagReadNoData is false
        if (!(flags & flagReadNoData))
        {
            std::memcpy(buf, data.data() + readOffset, readSize);
        }
        readOffset += readSize;

        return readSize;
    }

    static constexpr uint32_t flagWriteFail = 1 << 0;   // fail write with EIO
    static constexpr uint32_t flagWriteIntr = 1 << 1;   // fail write with EINTR
    static constexpr uint32_t flagWriteShort = 1 << 2;  // partial write
    static constexpr uint32_t flagWriteNoData = 1 << 3; // no data write

    static constexpr uint32_t flagReadFail = 1 << 4;   // fail read with EIO
    static constexpr uint32_t flagReadIntr = 1 << 5;   // fail read with EINTR
    static constexpr uint32_t flagReadShort = 1 << 6;  // partial read
    static constexpr uint32_t flagReadNoData = 1 << 7; // no data read

    static constexpr uint32_t flagCloseFail = 1 << 8; // fail close

    static constexpr auto deviceSize = 64 * 1024;
    static constexpr auto eraseSize = 1024;
    static constexpr auto intrCount = 3; // Number of times to fail with EINTR
    static constexpr auto shortCount =
        7; // Amount to decrease read/write size by

    mtd_info_t mtd = {};
    off_t writeOffset = 0, readOffset = 0;
    std::vector<uint8_t> data;
    uint32_t flags;
    int readIntrRemain = 0, writeIntrRemain = 0;
};

TEST_P(FlashCopyFakeTest, fakeTest)
{
    FakeMtd file(GetParam().first);
    std::vector<uint8_t> data(FakeMtd::deviceSize);
    std::mt19937 gen(0);
    std::generate(data.begin(), data.end(), gen);

    EXPECT_CALL(sys, open(StrEq(TEST_PATH), _)).WillOnce(Return(TEST_FD));
    EXPECT_CALL(sys, ioctl(TEST_FD, MEMGETINFO, NotNull()))
        .WillOnce(DoAll(SetArg2ToMtdPointer(file.mtd), Return(0)));
    EXPECT_CALL(
        sys, ioctl(TEST_FD, MEMERASE,
                   MatcherCast<void*>(MatcherCast<erase_info_t*>(
                       AllOf(Field(&erase_info_t::start, 0),
                             Field(&erase_info_t::length, file.deviceSize))))))
        .WillOnce(Return(0));
    EXPECT_CALL(sys, write(TEST_FD, NotNull(), _))
        .WillRepeatedly(Invoke(&file, &FakeMtd::write));
    EXPECT_CALL(sys, read(TEST_FD, NotNull(), _))
        .WillRepeatedly(Invoke(&file, &FakeMtd::read));

    // Mock sys->close call, allow error injection with flagCloseFail
    if (file.flags & FakeMtd::flagCloseFail)
    {
        EXPECT_CALL(sys, close(TEST_FD)).WillOnce(Return(-1));
    }
    else
    {
        EXPECT_CALL(sys, close(TEST_FD)).WillOnce(Return(0));
    }

    // This conditional checks whether we expect flashCopy to succeed
    if (GetParam().second)
    {
        google::hoth::internal::mtdImpl.flashCopy(&sys, data, TEST_PATH);

        EXPECT_THAT(file.data, ContainerEq(data));
    }
    else
    {
        EXPECT_THROW(
            google::hoth::internal::mtdImpl.flashCopy(&sys, data, TEST_PATH),
            std::exception);
    }
}

const std::vector<std::pair<uint32_t, bool>> fakeTestParams{
    // Params : {flags, expectation that flashCopy will succeed}

    // Passing case without any flags set
    {0, true},
    // Setting flags for sys->close mock
    {FakeMtd::flagCloseFail, false},
    // Setting flags for sys->write mock
    {FakeMtd::flagWriteIntr, true},
    {FakeMtd::flagWriteShort, true},
    {FakeMtd::flagWriteIntr | FakeMtd::flagWriteShort, true},
    {FakeMtd::flagWriteFail, false},
    {FakeMtd::flagWriteNoData, false},
    {FakeMtd::flagWriteIntr | FakeMtd::flagWriteShort | FakeMtd::flagWriteFail |
         FakeMtd::flagWriteNoData,
     false},
    // Setting flags for sys->read mock
    {FakeMtd::flagReadIntr, true},
    {FakeMtd::flagReadShort, true},
    {FakeMtd::flagReadIntr | FakeMtd::flagReadShort, true},
    {FakeMtd::flagReadFail, false},
    {FakeMtd::flagReadNoData, false},
    {FakeMtd::flagReadIntr | FakeMtd::flagReadShort | FakeMtd::flagReadFail |
         FakeMtd::flagReadNoData,
     false},
    // Corner case where we do not write or read data
    {FakeMtd::flagWriteNoData | FakeMtd::flagReadNoData, false}};

INSTANTIATE_TEST_CASE_P(WithFlags, FlashCopyFakeTest,
                        ::testing::ValuesIn(fakeTestParams));