// 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.hpp"

#include "fs.hpp"
#include "sys.hpp"

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

#include <algorithm>
#include <cstring>
#include <format>
#include <stdexcept>
#include <string>

namespace google
{
namespace hoth
{
namespace internal
{

// Sysfs directory containing info about mtd devices
auto constexpr sysfsMtd = "/sys/class/mtd";

std::string MtdImpl::findPartition(const std::string& name)
{
    auto dirs = listDir(sysfsMtd);

    for (const auto& p : dirs)
    {
        std::string line;

        try
        {
            line = getFirstLine(p / "name");
        }
        catch (const std::ios_base::failure& e)
        {
            // If we fail to read a name, just try the next one
            continue;
        }

        if (line == name)
        {
            return p.root_path() / "dev" / p.filename();
        }
    }

    throw std::runtime_error(
        std::format("Unable to find partition with name \"{}\"", name));
}

// TODO: choose the best block size
constexpr size_t blockSize = 4 * 1024;

void MtdImpl::flashCopy(const Sys* sys, const std::vector<uint8_t>& data,
                        const std::string& device)
{
    int fd = sys->open(device.c_str(), O_RDWR | O_SYNC);

    if (fd < 0)
    {
        throw errnoException(std::format("Failed to open device {}", device));
    }

    mtd_info_t mtd;
    if (sys->ioctl(fd, MEMGETINFO, &mtd) < 0)
    {
        sys->close(fd);

        throw errnoException(std::format("Failed to get info on {}", device));
    }

    if (data.size() > mtd.size)
    {
        sys->close(fd);

        throw std::runtime_error(std::format(
            "Data size {} too large for device {}", data.size(), device));
    }

    erase_info_t erase;
    erase.start = 0;
    erase.length = mtd.size;
    if (sys->ioctl(fd, MEMERASE, &erase) < 0)
    {
        sys->close(fd);

        throw errnoException(
            std::format("Failed to erase mtd device {}", device));
    }

    size_t count = 0;
    while (count < data.size())
    {
        size_t writeSize = std::min(data.size() - count, blockSize);
        auto ret = sys->write(fd, data.data() + count, writeSize);

        if (ret < 0)
        {
            if (errno == EINTR)
            {
                continue; // interrupted; try again
            }
            sys->close(fd);

            throw errnoException(
                std::format("Failed to write to mtd device {}", device));
        }
        if (ret == 0) 
        {
            sys->close(fd);

            throw std::runtime_error(std::format(
                "Unexpected EOF while writing to mtd device {}", device));
        } // ret > 0
        count += ret;
    }

    // Validate written data
    std::vector<uint8_t> readBuffer(blockSize);

    // Reset the file descriptor offset back to 0 before reading
    sys->lseek(fd, 0, SEEK_SET);

    count = 0;
    while (count < data.size())
    {
        size_t readSize = std::min(data.size() - count, blockSize);
        auto ret = sys->read(fd, readBuffer.data(), readSize);

        if (ret < 0)
        {
            if (errno == EINTR)
            {
                continue; // interrupted; try again
            }
            sys->close(fd);

            throw errnoException(
                std::format("Failed to read from mtd device {}", device));
        }
        if (ret == 0)
        {
            sys->close(fd);

            throw std::runtime_error(std::format(
                "Unexpected EOF while reading from mtd device {}", device));
        } // ret > 0
        int rc = std::memcmp(data.data() + count, readBuffer.data(), ret);
        if (rc != 0)
        {
            sys->close(fd);

            throw std::runtime_error(
                std::format("Written data failed validation at {} bytes", count));
        }
        count += ret;
    }

    auto retClose = sys->close(fd);
    if (retClose < 0)
    {
        throw errnoException(std::format("Failed to close device {}", device));
    }
}

MtdImpl mtdImpl;

} // namespace internal

} // namespace hoth

} // namespace google