//
// Copyright (C) 2023 The Android Open Source Project
//
// 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 "task.h"

#include <android-base/file.h>
#include <android-base/parseint.h>
#include <android-base/unique_fd.h>

#include "fastboot.h"
#include "fastboot_driver_interface.h"
#include "filesystem.h"
#include "liblp/builder.h"
#include "liblp/liblp.h"
#include "liblp/metadata_format.h"
#include "sparse/sparse.h"
#include "super_flash_helper.h"
#include "util.h"

// Need to include logging.h after fastboot.h, due to ERROR macro conflict
// clang-format off
#include <android-base/logging.h>
// clang-format on

using namespace std::string_literals;
using android::fs_mgr::GetPrimaryMetadataOffset;
using android::fs_mgr::MetadataBuilder;
using android::fs_mgr::ParseSuperPartition;
using android::fs_mgr::SlotNumberForSlotSuffix;
using android::fs_mgr::ValidateAndSerializeMetadata;

FlashTask::FlashTask(const std::string& slot, const std::string& pname, const std::string& fname,
                     const bool apply_vbmeta, const FlashingPlan* fp)
    : pname_(pname), fname_(fname), slot_(slot), apply_vbmeta_(apply_vbmeta), fp_(fp) {}

bool FlashTask::IsDynamicPartition(const ImageSource* source, const FlashTask* task) {
    std::vector<char> contents;
    if (!source->ReadFile("super_empty.img", &contents)) {
        return false;
    }
    auto metadata = android::fs_mgr::ReadFromImageBlob(contents.data(), contents.size());
    return should_flash_in_userspace(*metadata.get(), task->GetPartitionAndSlot());
}

fastboot::RetCode FlashPartition(fastboot::IFastBootDriver* fb, const std::string& partition_name,
                                 void* data, size_t size, off64_t offset) {
    // We use LP_PARTITION_RESERVED_BYTES as block size because everything in
    // liblp is aligned with this granularity
    std::unique_ptr<struct sparse_file, decltype(&sparse_file_destroy)> sparse_file(
            sparse_file_new(LP_PARTITION_RESERVED_BYTES, offset + size), sparse_file_destroy);
    // Fastboot's flash protocol does not have a way to specify offset we want
    // to write to, so we use sparse file as a way around it. Blocks which
    // have no data will be skipped by bootloader, similar to seeking forward
    // on a file.
    sparse_file_add_data(sparse_file.get(), data, size, offset / LP_PARTITION_RESERVED_BYTES);
    auto ret = fb->Download(partition_name, sparse_file.get(),
                            sparse_file_len(sparse_file.get(), true, false), 1, 1, false);
    if (ret != fastboot::SUCCESS) {
        return ret;
    }
    return fb->Flash(partition_name);
}

fastboot::RetCode FlashMetadata(MetadataBuilder* builder, uint32_t slot_number,
                                fastboot::IFastBootDriver* fb) {
    auto metadata = builder->Export();
    auto serialized = ValidateAndSerializeMetadata(*metadata);
    return FlashPartition(fb, "super", serialized.data(), serialized.size(),
                          GetPrimaryMetadataOffset(metadata->geometry, slot_number));
}

std::unique_ptr<MetadataBuilder> FetchAndParseMetadata(fastboot::IFastBootDriver* fb,
                                                       uint32_t slot_number) {
    // For typical use cases(3 slot + 65536 max metadata size), 512K is enough
    // We reserve 1MB just to get some headroom
    constexpr size_t kMaxMetadataSize = 1024 * 1024;
    std::vector<uint8_t> buffer;
    buffer.reserve(kMaxMetadataSize);
    // So that fb->Fetch() commands do not just crash on failure
    fastboot::NoCrashGuard guard(fb);
    auto ret = fb->Fetch("super", &buffer, 0, buffer.capacity());
    if (ret != fastboot::SUCCESS) {
        LOG(ERROR) << "Failed to fetch first " << kMaxMetadataSize << " bytes of super partition";
        return nullptr;
    }
    LpMetadataGeometry geometry;
    if (!android::fs_mgr::ParseGeometryFromSuperPartition(buffer.data(), buffer.size(),
                                                          &geometry)) {
        return nullptr;
    }
    const size_t total_metadata_size = android::fs_mgr::GetTotalMetadataSize(
            geometry.metadata_max_size, geometry.metadata_slot_count);
    if (total_metadata_size > kMaxMetadataSize) {
        LOG(ERROR) << "Total metadata size " << total_metadata_size
                   << " exceeds maximum supported metadata size " << kMaxMetadataSize;
        return nullptr;
    }
    auto metadata = ParseSuperPartition(buffer.data(), buffer.size(), slot_number);
    if (metadata == nullptr) {
        return nullptr;
    }
    return MetadataBuilder::New(*metadata);
}

constexpr uint64_t SectorToBlock(uint64_t sectors) {
    constexpr auto kSectorPerBlock = android::fs_mgr::kDefaultBlockSize / LP_SECTOR_SIZE;
    return sectors / kSectorPerBlock;
}

android::base::unique_fd OpenFile(const FlashingPlan* fp, const char* fname) {
    if (fp->source) {
        return fp->source->OpenFile(fname);
    } else {
        return unique_fd(TEMP_FAILURE_RETRY(open(fname, O_RDONLY | O_BINARY)));
    }
}

// Generate a sparse image where extents from LpMetadata partition |p| is
// filled with data from |fname| . This sparse image can then be sent over
// fastboot protocol to flash super partition, end result is flashing a single
// partition inside super.
SparsePtr SparseForSuper(const char* fname, int64_t file_size, android::fs_mgr::Partition* p) {
    struct fastboot_buffer buf;

    SparsePtr sparse_file = SparsePtr(
            sparse_file_new(android::fs_mgr::kDefaultBlockSize, file_size), sparse_file_destroy);
    uint64_t offset = 0;
    for (const auto& extent : p->extents()) {
        auto e = extent->AsLinearExtent();
        if (!e) {
            die("Partition on device contains a non-linear extent. unsupported.");
        }
        sparse_file_add_file(sparse_file.get(), fname, offset,
                             extent->num_sectors() * LP_SECTOR_SIZE,
                             SectorToBlock(e->physical_sector()));
        offset += extent->num_sectors() * LP_SECTOR_SIZE;
    }

    LOG(INFO) << "Transfer size for " << p->name() << " is " << offset;
    return sparse_file;
}

void FlashTask::Run() {
    const auto is_fastbootd = is_userspace_fastboot(fp_->fb);
    const auto slot_number = SlotNumberForSlotSuffix(fp_->current_slot);

    auto flash = [&](const std::string& partition) {
        const bool is_dynamic_partition = should_flash_in_userspace(fp_->source.get(), partition);
        // fastbootd should be able to flash both static and dynamic
        // partitions, just flash it
        if (is_fastbootd) {
            do_flash(partition.c_str(), fname_.c_str(), apply_vbmeta_, fp_);
            return;
        }
        if (!is_dynamic_partition || fp_->force_flash) {
            do_flash(partition.c_str(), fname_.c_str(), apply_vbmeta_, fp_);
            return;
        }
        // is_dynamic && !force_flash
        // Make last attempt to save everything by fetching super partition
        // metadata
        auto builder = FetchAndParseMetadata(fp_->fb, slot_number);
        if (builder == nullptr) {
            die("The partition you are trying to flash is dynamic, and "
                "should be flashed via fastbootd. Please run:\n"
                "\n"
                "    fastboot reboot fastboot\n"
                "\n"
                "And try again. If you are intentionally trying to "
                "overwrite a fixed partition, use --force.");
        }
        auto part = builder->FindPartition(partition);
        if (part == nullptr) {
            die("The partition you are trying to flash is not a static partition, and we cannot "
                "find it in super partition either");
        }
        auto fd = OpenFile(fp_, fname_.c_str());
        if (fd < 0) {
            die("Failed to open input file %s : %s", fname_.c_str(), strerror(errno));
        }
        if (is_sparse_file(fd)) {
            die("Flashing sparse images to dynamic partitions is currently not supported");
        }
        auto file_size = get_file_size(fd);
        if (file_size <= 0) {
            die("Failed to get file size for %s : %s", fname_.c_str(), strerror(errno));
        }
        if ((uint64_t)file_size != part->size()) {
            part->RemoveExtents();
            builder->ResizePartition(part, file_size);
        }
        SparsePtr sparse_file = SparseForSuper(fname_.c_str(), file_size, part);

        auto metadata = builder->Export();
        auto serialized = ValidateAndSerializeMetadata(*metadata);
        sparse_file_add_data(sparse_file.get(), serialized.data(), serialized.size(),
                             GetPrimaryMetadataOffset(metadata->geometry, slot_number) /
                                     android::fs_mgr::kDefaultBlockSize);

        flash_partition(fp_, "super", std::move(sparse_file));
    };
    do_for_partitions(fp_->fb, pname_, slot_, flash, true);
}

std::string FlashTask::ToString() const {
    std::string apply_vbmeta_string = "";
    if (apply_vbmeta_) {
        apply_vbmeta_string = " --apply_vbmeta";
    }
    return "flash" + apply_vbmeta_string + " " + pname_ + " " + fname_;
}

std::string FlashTask::GetPartitionAndSlot() const {
    auto slot = slot_;
    if (slot.empty()) {
        slot = get_current_slot(fp_->fb);
    }
    if (slot.empty()) {
        return pname_;
    }
    if (slot == "all") {
        LOG(FATAL) << "Cannot retrieve a singular name when using all slots";
    }
    return pname_ + "_" + slot;
}

RebootTask::RebootTask(const FlashingPlan* fp) : fp_(fp){};
RebootTask::RebootTask(const FlashingPlan* fp, const std::string& reboot_target)
    : reboot_target_(reboot_target), fp_(fp){};

void RebootTask::Run() {
    if (reboot_target_ == "fastboot") {
        if (!is_userspace_fastboot(fp_->fb)) {
            reboot_to_userspace_fastboot(fp_->fb);
            fp_->fb->WaitForDisconnect();
        }
    } else if (reboot_target_ == "recovery") {
        fp_->fb->RebootTo("recovery");
        fp_->fb->WaitForDisconnect();
    } else if (reboot_target_ == "bootloader") {
        fp_->fb->RebootTo("bootloader");
        fp_->fb->WaitForDisconnect();
    } else if (reboot_target_ == "") {
        fp_->fb->Reboot();
        fp_->fb->WaitForDisconnect();
    } else {
        syntax_error("unknown reboot target %s", reboot_target_.c_str());
    }
}

std::string RebootTask::ToString() const {
    return "reboot " + reboot_target_;
}

OptimizedFlashSuperTask::OptimizedFlashSuperTask(const std::string& super_name,
                                                 std::unique_ptr<SuperFlashHelper> helper,
                                                 SparsePtr sparse_layout, uint64_t super_size,
                                                 const FlashingPlan* fp)
    : super_name_(super_name),
      helper_(std::move(helper)),
      sparse_layout_(std::move(sparse_layout)),
      super_size_(super_size),
      fp_(fp) {}

void OptimizedFlashSuperTask::Run() {
    // Use the reported super partition size as the upper limit, rather than
    // sparse_file_len, which (1) can fail and (2) is kind of expensive, since
    // it will map in all of the embedded fds.
    std::vector<SparsePtr> files;
    if (int limit = get_sparse_limit(super_size_, fp_)) {
        if (!split_file(sparse_layout_.get(), limit, &files)) {
            LOG(FATAL) << "Failed to resparse super partition";
        }
    } else {
        files.emplace_back(std::move(sparse_layout_));
    }

    // Send the data to the device.
    flash_partition_files(fp_->fb, super_name_, files);
}

std::string OptimizedFlashSuperTask::ToString() const {
    return "optimized-flash-super";
}

// This looks for a block within tasks that has the following pattern [reboot fastboot,
// update-super, $LIST_OF_DYNAMIC_FLASH_TASKS] and returns true if this is found.Theoretically
// this check is just a pattern match and could break if fastboot-info has a bunch of junk commands
// but all devices should pretty much follow this pattern
bool OptimizedFlashSuperTask::CanOptimize(const ImageSource* source,
                                          const std::vector<std::unique_ptr<Task>>& tasks) {
    for (size_t i = 0; i < tasks.size(); i++) {
        auto reboot_task = tasks[i]->AsRebootTask();
        if (!reboot_task || reboot_task->GetTarget() != "fastboot") {
            continue;
        }
        // The check for i >= tasks.size() - 2 is because we are peeking two tasks ahead. We need to
        // check for an update-super && flash {dynamic_partition}
        if (i >= tasks.size() - 2 || !tasks[i + 1]->AsUpdateSuperTask()) {
            continue;
        }
        auto flash_task = tasks[i + 2]->AsFlashTask();
        if (!FlashTask::IsDynamicPartition(source, flash_task)) {
            continue;
        }
        return true;
    }
    return false;
}

std::unique_ptr<OptimizedFlashSuperTask> OptimizedFlashSuperTask::Initialize(
        const FlashingPlan* fp, std::vector<std::unique_ptr<Task>>& tasks) {
    if (!fp->should_optimize_flash_super) {
        LOG(INFO) << "super optimization is disabled";
        return nullptr;
    }
    if (!supports_AB(fp->fb)) {
        LOG(VERBOSE) << "Cannot optimize flashing super on non-AB device";
        return nullptr;
    }
    if (fp->slot_override == "all") {
        LOG(VERBOSE) << "Cannot optimize flashing super for all slots";
        return nullptr;
    }
    if (!CanOptimize(fp->source.get(), tasks)) {
        return nullptr;
    }

    // Does this device use dynamic partitions at all?
    unique_fd fd = fp->source->OpenFile("super_empty.img");

    if (fd < 0) {
        LOG(VERBOSE) << "could not open super_empty.img";
        return nullptr;
    }

    std::string super_name;
    // Try to find whether there is a super partition.
    if (fp->fb->GetVar("super-partition-name", &super_name) != fastboot::SUCCESS) {
        super_name = "super";
    }
    uint64_t partition_size;
    std::string partition_size_str;
    if (fp->fb->GetVar("partition-size:" + super_name, &partition_size_str) != fastboot::SUCCESS) {
        LOG(VERBOSE) << "Cannot optimize super flashing: could not determine super partition";
        return nullptr;
    }
    partition_size_str = fb_fix_numeric_var(partition_size_str);
    if (!android::base::ParseUint(partition_size_str, &partition_size)) {
        LOG(VERBOSE) << "Could not parse " << super_name << " size: " << partition_size_str;
        return nullptr;
    }

    std::unique_ptr<SuperFlashHelper> helper = std::make_unique<SuperFlashHelper>(*fp->source);
    if (!helper->Open(fd)) {
        return nullptr;
    }

    for (const auto& task : tasks) {
        if (auto flash_task = task->AsFlashTask()) {
            auto partition = flash_task->GetPartitionAndSlot();
            if (!helper->AddPartition(partition, flash_task->GetImageName(), false)) {
                return nullptr;
            }
        }
    }

    auto s = helper->GetSparseLayout();
    if (!s) return nullptr;

    // Remove tasks that are concatenated into this optimized task
    auto remove_if_callback = [&](const auto& task) -> bool {
        if (auto flash_task = task->AsFlashTask()) {
            return helper->WillFlash(flash_task->GetPartitionAndSlot());
        } else if (task->AsUpdateSuperTask()) {
            return true;
        } else if (auto reboot_task = task->AsRebootTask()) {
            if (reboot_task->GetTarget() == "fastboot") {
                return true;
            }
        }
        return false;
    };

    tasks.erase(std::remove_if(tasks.begin(), tasks.end(), remove_if_callback), tasks.end());

    return std::make_unique<OptimizedFlashSuperTask>(super_name, std::move(helper), std::move(s),
                                                     partition_size, fp);
}

UpdateSuperTask::UpdateSuperTask(const FlashingPlan* fp) : fp_(fp) {}

void UpdateSuperTask::Run() {
    unique_fd fd = fp_->source->OpenFile("super_empty.img");
    if (fd < 0) {
        return;
    }
    if (!is_userspace_fastboot(fp_->fb)) {
        reboot_to_userspace_fastboot(fp_->fb);
    }

    std::string super_name;
    if (fp_->fb->GetVar("super-partition-name", &super_name) != fastboot::RetCode::SUCCESS) {
        super_name = "super";
    }
    fp_->fb->Download(super_name, fd, get_file_size(fd));

    std::string command = "update-super:" + super_name;
    if (fp_->wants_wipe) {
        command += ":wipe";
    }
    fp_->fb->RawCommand(command, "Updating super partition");
}
std::string UpdateSuperTask::ToString() const {
    return "update-super";
}

ResizeTask::ResizeTask(const FlashingPlan* fp, const std::string& pname, const std::string& size,
                       const std::string& slot)
    : fp_(fp), pname_(pname), size_(size), slot_(slot) {}

void ResizeTask::Run() {
    if (is_userspace_fastboot(fp_->fb)) {
        auto resize_partition = [this](const std::string& partition) -> void {
            if (is_logical(fp_->fb, partition)) {
                fp_->fb->ResizePartition(partition, size_);
            }
        };
        do_for_partitions(fp_->fb, pname_, slot_, resize_partition, false);
        return;
    }
    auto slot_number = SlotNumberForSlotSuffix(fp_->current_slot);
    auto builder = FetchAndParseMetadata(fp_->fb, slot_number);
    auto resize_partition = [this,
                             builder(builder.get())](const std::string& partition) mutable -> void {
        auto p = builder->FindPartition(partition);
        if (p) {
            builder->ResizePartition(p, atoll(size_.c_str()));
        }
    };
    do_for_partitions(fp_->fb, pname_, slot_, resize_partition, false);
    FlashMetadata(builder.get(), slot_number, fp_->fb);
}

std::string ResizeTask::ToString() const {
    return "resize " + pname_;
}

DeleteTask::DeleteTask(const FlashingPlan* fp, const std::string& pname) : fp_(fp), pname_(pname){};

void DeleteTask::Run() {
    fp_->fb->DeletePartition(pname_);
}

std::string DeleteTask::ToString() const {
    return "delete " + pname_;
}

WipeTask::WipeTask(const FlashingPlan* fp, const std::string& pname) : fp_(fp), pname_(pname){};

void WipeTask::Run() {
    std::string partition_type;
    if (fp_->fb->GetVar("partition-type:" + pname_, &partition_type) != fastboot::SUCCESS) {
        LOG(ERROR) << "wipe task partition not found: " << pname_;
        return;
    }
    if (partition_type.empty()) return;
    if (fp_->fb->Erase(pname_) != fastboot::SUCCESS) {
        LOG(ERROR) << "wipe task erase failed with partition: " << pname_;
        return;
    }
    fb_perform_format(pname_, 1, partition_type, "", fp_->fs_options, fp_);
}

std::string WipeTask::ToString() const {
    return "erase " + pname_;
}
