// Copyright 2021 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 <flasher/logging.hpp>
#include <flasher/ops.hpp>

#include <algorithm>
#include <cstddef>
#include <format>
#include <stdexcept>
#include <vector>

namespace flasher
{
namespace ops
{

static void readAtExact(Device& dev, std::span<std::byte> buf, size_t off)
{
    dev.readAtExact(buf, off);
    LOG(LogLevel::Info, " RD@{}#{}", off, buf.size());
}

static void writeAtExact(Device& dev, std::span<const std::byte> data,
                         size_t off)
{
    dev.writeAtExact(data, off);
    LOG(LogLevel::Info, " WD@{}#{}", off, data.size());
}

static void eraseBlocks(Device& dev, size_t off, size_t size)
{
    dev.eraseBlocks(off, size);
    LOG(LogLevel::Info, " ED@{}#{}", dev.getEraseSize() * off,
        dev.getEraseSize() * size);
}

void erase(Device& dev, size_t dev_offset, size_t max_size,
           std::optional<size_t> stride_size, bool noread)
{
    size_t stride = stride_size ? *stride_size : dev.recommendedStride();
    if (stride == 0)
    {
        throw std::invalid_argument("Stride cannot be 0");
    }
    stride = dev.eraseAlignEnd(stride);
    if (dev_offset > dev.getSize())
    {
        throw std::invalid_argument(std::format(
            "Device smaller than offset, {} < {}", dev.getSize(), dev_offset));
    }
    max_size = std::min(max_size, dev.getSize() - dev_offset);
    if (max_size == 0)
    {
        LOG(LogLevel::Info, "\n");
        return;
    }
    std::vector<std::byte> buf_v(stride * 2), erased_v(stride * 2);
    dev.mockErase(erased_v);
    const std::span<std::byte> buf(buf_v), erased(erased_v);
    const auto erase_size = dev.getEraseSize();
    const auto start = dev.eraseAlignStart(dev_offset);
    if (start != dev_offset)
    {
        const auto start_size = dev_offset - start;
        max_size += start_size;
        auto erase = true;
        if (noread)
        {
            readAtExact(dev, buf.subspan(0, start_size), start);
            if (max_size < erase_size)
            {
                readAtExact(dev, buf.subspan(max_size, erase_size - max_size),
                            start + max_size);
            }
        }
        else
        {
            readAtExact(dev, buf.subspan(0, erase_size), start);
            const auto cmp_size = std::min(max_size, erase_size) - start_size;
            if (!dev.needsErase(buf.subspan(start_size, cmp_size),
                                erased.subspan(start_size, cmp_size)))
            {
                erase = false;
            }
        }
        if (erase)
        {
            eraseBlocks(dev, start / erase_size, 1);
            writeAtExact(dev, buf.subspan(0, start_size), start);
            if (max_size < erase_size)
            {
                writeAtExact(dev, buf.subspan(max_size, erase_size - max_size),
                             start + max_size);
            }
        }
        if (max_size < erase_size)
        {
            LOG(LogLevel::Info, "\n");
            return;
        }
        dev_offset = start + erase_size;
        max_size -= erase_size;
    }
    const auto uend = dev_offset + max_size;
    const auto end = dev.eraseAlignEnd(uend);
    if (end != uend)
    {
        const auto end_size = end - uend;
        auto erase = true;
        if (noread)
        {
            readAtExact(dev, buf.subspan(erase_size - end_size, end_size),
                        uend);
        }
        else
        {
            readAtExact(dev, buf.subspan(0, erase_size), end - erase_size);
            if (!dev.needsErase(buf.subspan(0, erase_size - end_size),
                                erased.subspan(0, erase_size - end_size)))
            {
                erase = false;
            }
        }
        if (erase)
        {
            eraseBlocks(dev, end / erase_size - 1, 1);
            writeAtExact(dev, buf.subspan(erase_size - end_size, end_size),
                         uend);
        }
        max_size -= erase_size - end_size;
    }
    for (size_t i = dev_offset; i < max_size + dev_offset; i += stride)
    {
        const size_t bytes = std::min(stride, max_size + dev_offset - i);
        if (!noread)
        {
            readAtExact(dev, buf.subspan(0, bytes), i);
            if (!dev.needsErase(buf.subspan(0, bytes),
                                erased.subspan(0, bytes)))
            {
                continue;
            }
        }
        eraseBlocks(dev, i / erase_size, bytes / erase_size);
    }
    LOG(LogLevel::Info, "\n");
}

} // namespace ops
} // namespace flasher