// 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/file.hpp>
#include <flasher/file/memory.hpp>
#include <flasher/mod.hpp>
#include <flasher/ops.hpp>
#include <flashupdate/args.hpp>
#include <flashupdate/flash.hpp>
#include <flashupdate/info.hpp>
#include <flashupdate/logging.hpp>
#include <flashupdate/ops.hpp>
#include <flashupdate/validator.hpp>
#include <flashupdate/version.hpp>

#include <filesystem>
#include <format>
#include <string>

namespace flashupdate
{
namespace ops
{

using flasher::ModArgs;
using stdplus::fd::OpenAccess;
using stdplus::fd::OpenFlag;
using stdplus::fd::OpenFlags;

inline void
    writeWithoutValidation(flash::Flash* flashHelper,
                           flasher::Reader& inputImage,
                           std::span<std::pair<size_t, size_t>> partitions)
{
    // Secondary partition only.
    auto flash = flashHelper->getFlash(false, inputImage.getSize());
    if (!flash)
    {
        throw std::runtime_error(
            std::format("failed to find secondary Flash partition"));
    }

    // Validate the image currently in the flash.
    // It can be the primary or secondary flash.
    std::string flashDev(flash->first);
    auto devMod = ModArgs(flashDev);
    flashDev = devMod.arr.back();
    flasher::NestedMutate mutate{};
    auto dev = flasher::openDevice(devMod);
    LOG(LogLevel::Info, "Flash image to {}\n", flashDev);
    for (auto& [offset, size] : partitions)
    {
        flasher::ops::automatic(/*dst=*/*dev, /*dst_offset=*/offset,
                                /*src=*/inputImage, /*src_offset=*/offset,
                                mutate, /*max_size=*/size, std::nullopt,
                                /*noread=*/false);
    }
}

void writeHelper(const Args& args, std::string_view image,
                 flasher::Reader& inputImage)
{
    std::vector<std::pair<size_t, size_t>> partitions;
    size_t size = inputImage.getSize();

    // Write without validation only if we are skipping validation and not
    // updating the staging index. The staging index is the index that points to
    // the partition that is allowed to write to the primary partition and
    // validation is required to update it.
    if (args.skipValidation && !args.updateStagingIndex)
    {
        if (args.primary)
        {
            throw std::runtime_error(
                std::format("Cannot skip validation on primary partition"));
        }

        args.flashHelper->setup(args.config, args.keepMux);
        partitions.emplace_back(0, size);
        writeWithoutValidation(args.flashHelper, inputImage, partitions);
        return;
    }

    auto& helper = args.validatorHelper;
    if (helper == nullptr)
    {
        throw std::runtime_error("invalid Validator Helper");
    }

    if (!helper->validateImage(inputImage, args.config.flash.validationKey))
    {
        throw std::runtime_error(std::format(
            "failed to validate the CR51 descriptor for the next image: {}",
            image));
    }

    if (args.config.supportedVersion != std::nullopt)
    {
        std::string versionStr = helper->imageVersion();
        auto version = version::Version(versionStr);

        bool validVersion = true;
        for (const auto& supportedVersion : *args.config.supportedVersion)
        {
            if (
                // Less than min version
                (supportedVersion.min != std::nullopt &&
                 version < *supportedVersion.min) ||
                // greater than or equal to max version
                (supportedVersion.max != std::nullopt &&
                 version >= *supportedVersion.max))
            {
                validVersion = false;
                break;
            }
        }

        if (!validVersion)
        {
            throw std::runtime_error(std::format(
                "Next image's version is not supported for {}", versionStr));
        }
    }

    auto& flashHelper = args.flashHelper;
    flashHelper->setup(args.config, args.keepMux);

    auto flash = flashHelper->getFlash(args.primary, size);
    if (!flash)
    {
        throw std::runtime_error("failed to find Flash partition");
    }

    // Save the information from the next image
    auto nextDescriptorHashSpan = helper->descriptorHash().hash;
    std::vector<uint8_t> nextDescriptorHash(nextDescriptorHashSpan.begin(),
                                            nextDescriptorHashSpan.end());
    auto nextIsProd = helper->prodImage();

    // Validate the image currently in the flash.
    // It can be the primary or secondary flash.
    std::string flashDev(flash->first);
    auto devMod = ModArgs(flashDev);
    flashDev = devMod.arr.back();

    // Fetch BIOS Staging Information
    auto info = fetchInfo(args);

    // Check Image Descriptor Hash
    // Prevent flashing the image to primary partition if the hash of image
    // descriptor does not match the cached.
    // If the primary partition is invalid, skip the check and allow firmware
    // update.
    if (args.primary)
    {
        auto metadataDescriptor = info.stageDescriptor;
        std::string flashDev(flash->first);

        bool currentIsProd = args.validatorHelper->isProdImage(
            devMod, metadataDescriptor.offset, metadataDescriptor.size);

        // Flash Support only for
        // Prod <-> Prod
        // Dev   -> Prod
        // Dev  <-> Dev
        // The Prod to dev is controlled through the image validation config.
        if (!args.config.cr51->prodToDev && currentIsProd && !nextIsProd)
        {
            throw std::logic_error("Prod to dev update is not enabled.");
        }

        if (args.stagingIndex != info.stagingIndex)
        {
            throw std::logic_error(
                std::format("The Staged Partition is not in the expected "
                            "partition: want {}, got {}",
                            info.stagingIndex, args.stagingIndex));
        }

        // Get the data to match the hash size in the installing image.
        std::span<const uint8_t> expectedHash(
            &(info.stageDescriptor.hash[0]),
            &(info.stageDescriptor.hash[0]) +
                std::min(static_cast<size_t>(SHA256_DIGEST_LENGTH),
                         nextDescriptorHash.size()));

        LOG(LogLevel::Info, "Checking HASH for CR51 descriptor between "
                            "staged partition and cache");

        if (!std::equal(nextDescriptorHash.begin(), nextDescriptorHash.end(),
                        expectedHash.begin(), expectedHash.end()))
        {
            throw std::logic_error(
                "SHA256 of the staged image in the cache "
                "does not match the image in the staged partition");
        }

        // Insert the non-persistent regions
        auto regions = helper->persistentRegions();
        size_t start = 0;
        for (const auto& region : regions)
        {
            LOG(LogLevel::Notice, "Region: {}, Inject offset: {}, Length: {}\n",
                region.name, region.offset, region.size);
            if (start == region.offset)
            {
                start += region.size;
                continue;
            }
            partitions.emplace_back(start, region.offset - start);
            start = region.offset + region.size;
        }
        partitions.emplace_back(start, size - start);
    }
    // Not writing to active/primary partition. Write the whole image
    else
    {
        partitions.emplace_back(0, size);
    }

    LOG(LogLevel::Info, "finished setting up {} and {}", image, flashDev);
    {
        // All checks passed, continue to flash the image
        flasher::NestedMutate mutate{};
        auto dev = flasher::openDevice(devMod);
        LOG(LogLevel::Info, "Flash image to {}", flashDev);
        for (auto& [offset, size] : partitions)
        {
            LOG(LogLevel::Debug, "Writing partition: offset {}, size {}",
                offset, size);
            flasher::ops::automatic(/*dst=*/*dev, /*dst_offset=*/offset,
                                    /*src=*/inputImage, /*src_offset=*/offset,
                                    mutate, /*max_size=*/size, std::nullopt,
                                    /*noread=*/false);
        }

        LOG(LogLevel::Info, "finished flashing {} to {}", image, flashDev);
        // Validate the flash after writing the image
        auto fileMod = ModArgs(flashDev);
        auto nextImage =
            flasher::openFile(fileMod, OpenFlags(OpenAccess::ReadOnly));
        if (!helper->validateImage(*nextImage, args.config.flash.validationKey))
        {
            throw std::runtime_error(std::format(
                "failed to validate the CR51 descriptor for the image "
                "in the flash after overwriting it: {}",
                image));
        }
    }
    info::UpdateInfo::State expectedState;

    // Update the Active Version the and State to UPDATED
    if (args.primary)
    {
        info.active = version::Version(helper->imageVersion());
        expectedState = info::UpdateInfo::State::UPDATED;
        info.state = static_cast<uint8_t>(expectedState);
    }
    else
    // Update the Staged version the the State to STAGED.
    //
    // Flashing to secondary flash will save the has of CR51 image descriptor
    // to the EEPROM for the final check when write to primary flash.
    {
        info.stage = version::Version(helper->imageVersion());
        expectedState = info::UpdateInfo::State::STAGED;
        info.state = static_cast<uint8_t>(expectedState);

        // Update the Staging Index to make sure the it uses the same staged
        // partition between the stage and active.
        if (args.updateStagingIndex)
        {
            info.stagingIndex = args.stagingIndex;
        }

        // Save the CR51 descriptor hash for the staged image
        // The hash size could be 16, 32, or 64 bytes depending on the Hash
        // type. Save up to 32 bytes (SHA256_DIGEST_LENGTH) of data.
        auto descriptor = helper->descriptorHash();
        auto descriptorHashSize = std::min(
            static_cast<size_t>(SHA256_DIGEST_LENGTH), descriptor.hash.size());
        memcpy(info.stageDescriptor.hash, descriptor.hash.data(),
               descriptorHashSize);
        info.stageDescriptor.offset = descriptor.offset;
        info.stageDescriptor.size = descriptor.size;
    }

    info::printUpdateInfo(args, info);

    auto findState = info::stateToString.find(info.state);
    if (findState != info::stateToString.end())
    {
        LOG(LogLevel::Notice, "Updated the staged BIOS to {}",
            findState->second);
    }
    writeInfo(args, info);

    // Confirm the Update State is what we expected
    info = fetchInfo(args);
    if (info.state != static_cast<uint8_t>(expectedState))
    {
        throw std::logic_error(std::format(
            "the update state is not what we expected: want {}, got {}",
            static_cast<uint8_t>(expectedState), info.state));
    }
}

void write(const Args& args)
{
    // Validate the next image before attempting to flash it
    std::string image = args.file->arr.back();
    LOG(LogLevel::Notice, "Validate BIOS image: {}", image);

    auto fileMod = ModArgs(image);
    auto inputImage =
        flasher::openFile(fileMod, OpenFlags(OpenAccess::ReadOnly));

    writeHelper(args, image, *inputImage);
}

} // namespace ops
} // namespace flashupdate