// 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/device.hpp>
#include <flasher/modutil.hpp>
#include <flasher/util.hpp>

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

namespace flasher
{

void Device::writeAtExact(std::span<const std::byte> data, size_t offset)
{
    opAtExact("Device writeAtExact", &Device::writeAt, *this, data, offset);
}

Device::Type Device::toType(std::string_view str)
{
    if (str == "nor")
    {
        return Type::Nor;
    }
    if (str == "simple")
    {
        return Type::Simple;
    }
    throw std::invalid_argument(std::format("Not a device type: {}", str));
}

Device::Device(DeviceInfo&& info) : info(std::move(info))
{
    if (info.erase_size != 0 && info.size % info.erase_size != 0)
    {
        throw std::invalid_argument(
            std::format("Flash size {} is not divisible by erase size {}",
                        info.size, info.erase_size));
    }
    switch (info.type)
    {
        case Type::Nor:
            if (info.erase_size == 0)
            {
                throw std::invalid_argument(
                    "Nor flashes can't have 0 erase size");
            }
            break;
        case Type::Simple:
            if (info.erase_size != 0)
            {
                throw std::invalid_argument("Simple flashes can't erase");
            }
            break;
    }
}

size_t Device::recommendedStride() const
{
    return info.erase_size == 0 ? 8192 : info.erase_size;
}

size_t Device::getSize() const
{
    return info.size;
}

size_t Device::eraseAlignStart(size_t offset) const
{
    if (info.erase_size == 0)
        return offset;
    return offset / info.erase_size * info.erase_size;
}

size_t Device::eraseAlignEnd(size_t offset) const
{
    if (info.erase_size == 0)
        return offset;
    return eraseAlignStart(offset + info.erase_size - 1);
}

bool Device::needsErase(std::span<const std::byte> flash_data,
                        std::span<const std::byte> new_data) const
{
    if (new_data.size() > flash_data.size())
    {
        throw std::invalid_argument("New data bigger than MTD data");
    }
    switch (info.type)
    {
        case Type::Nor:
            for (size_t i = 0; i < new_data.size(); ++i)
            {
                if ((flash_data[i] & new_data[i]) != new_data[i])
                {
                    return true;
                }
            }
            return false;
        case Type::Simple:
            return false;
    }
    throw std::invalid_argument("Don't know how to erase Device");
}

template <typename Iter>
static size_t shrinkWrite(Device::Type type, Iter flash_data_begin,
                          Iter new_data_begin, Iter new_data_end)
{
    switch (type)
    {
        case Device::Type::Nor:
            for (auto itn = new_data_begin, itf = flash_data_begin;
                 new_data_end != itn; itn++, itf++)
            {
                if ((*itn & *itf) != *itf)
                {
                    return itn - new_data_begin;
                }
            }
            return new_data_end - new_data_begin;
        case Device::Type::Simple:
            for (auto itn = new_data_begin, itf = flash_data_begin;
                 new_data_end != itn; itn++, itf++)
            {
                if (*itn != *itf)
                {
                    return itn - new_data_begin;
                }
            }
            return new_data_end - new_data_begin;
    }
    throw std::invalid_argument("Don't know how to shrink write data");
}

size_t Device::shrinkWritePre(std::span<const std::byte> flash_data,
                              std::span<const std::byte> new_data) const
{
    return shrinkWrite(info.type, flash_data.begin(), new_data.begin(),
                       new_data.begin() +
                           std::min(flash_data.size(), new_data.size()));
}

size_t Device::shrinkWritePost(std::span<const std::byte> flash_data,
                               std::span<const std::byte> new_data) const
{
    auto max = std::min(flash_data.size(), new_data.size());
    return max -
           shrinkWrite(info.type, flash_data.rbegin() + flash_data.size() - max,
                       new_data.rbegin() + new_data.size() - max,
                       new_data.rend());
}

void Device::mockErase(std::span<std::byte> data) const
{
    switch (info.type)
    {
        case Type::Nor:
            memset(data.data(), 0xff, data.size());
            return;
        case Type::Simple:
            return;
    }
    throw std::invalid_argument("Don't know how to erase Device");
}

ModTypeMap<DeviceType>& getDeviceTypes() noexcept
{
    static ModTypeMap<DeviceType> deviceTypes;
    return deviceTypes;
}

std::unique_ptr<Device> openDevice(const ModArgs& args)
{
    return openMod<Device>(getDeviceTypes(), args);
}

void registerDeviceType(std::string_view name,
                        std::unique_ptr<DeviceType>&& type) noexcept
{
    getDeviceTypes().emplace(name, std::move(type));
}

void unregisterDeviceType(std::string_view name) noexcept
{
    getDeviceTypes().erase(name);
}

} // namespace flasher