/*
 * Copyright (C) 2019 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 "host/libs/config/data_image.h"

#include <android-base/logging.h>
#include <android-base/result.h>

#include "blkid.h"

#include "common/libs/fs/shared_buf.h"
#include "common/libs/utils/files.h"
#include "common/libs/utils/result.h"
#include "common/libs/utils/subprocess.h"
#include "host/libs/config/cuttlefish_config.h"
#include "host/libs/config/esp.h"
#include "host/libs/config/mbr.h"
#include "host/libs/config/openwrt_args.h"
#include "host/libs/vm_manager/gem5_manager.h"

namespace cuttlefish {

using APBootFlow = CuttlefishConfig::InstanceSpecific::APBootFlow;
using BootFlow = CuttlefishConfig::InstanceSpecific::BootFlow;

namespace {

static constexpr std::string_view kDataPolicyUseExisting = "use_existing";
static constexpr std::string_view kDataPolicyAlwaysCreate = "always_create";
static constexpr std::string_view kDataPolicyResizeUpTo = "resize_up_to";

const int FSCK_ERROR_CORRECTED = 1;
const int FSCK_ERROR_CORRECTED_REQUIRES_REBOOT = 2;

Result<void> ForceFsckImage(
    const std::string& data_image,
    const CuttlefishConfig::InstanceSpecific& instance) {
  std::string fsck_path;
  if (instance.userdata_format() == "f2fs") {
    fsck_path = HostBinaryPath("fsck.f2fs");
  } else if (instance.userdata_format() == "ext4") {
    fsck_path = "/sbin/e2fsck";
  }
  int fsck_status = Execute({fsck_path, "-y", "-f", data_image});
  CF_EXPECTF(!(fsck_status &
               ~(FSCK_ERROR_CORRECTED | FSCK_ERROR_CORRECTED_REQUIRES_REBOOT)),
             "`{} -y -f {}` failed with code {}", fsck_path, data_image,
             fsck_status);
  return {};
}

Result<void> ResizeImage(const std::string& data_image, int data_image_mb,
                         const CuttlefishConfig::InstanceSpecific& instance) {
  auto file_mb = FileSize(data_image) >> 20;
  CF_EXPECTF(data_image_mb >= file_mb, "'{}' is already {} MB, won't downsize",
             data_image, file_mb);
  if (file_mb == data_image_mb) {
    LOG(INFO) << data_image << " is already the right size";
    return {};
  }
  off_t raw_target = static_cast<off_t>(data_image_mb) << 20;
  auto fd = SharedFD::Open(data_image, O_RDWR);
  CF_EXPECTF(fd->IsOpen(), "Can't open '{}': '{}'", data_image, fd->StrError());
  CF_EXPECTF(fd->Truncate(raw_target) == 0, "`truncate --size={}M {} fail: {}",
             data_image_mb, data_image, fd->StrError());
  CF_EXPECT(ForceFsckImage(data_image, instance));
  std::string resize_path;
  if (instance.userdata_format() == "f2fs") {
    resize_path = HostBinaryPath("resize.f2fs");
  } else if (instance.userdata_format() == "ext4") {
    resize_path = "/sbin/resize2fs";
  }
  if (resize_path != "") {
    CF_EXPECT_EQ(Execute({resize_path, data_image}), 0,
                 "`" << resize_path << " " << data_image << "` failed");
    CF_EXPECT(ForceFsckImage(data_image, instance));
  }

  return {};
}

std::string GetFsType(const std::string& path) {
  std::string fs_type;
  blkid_cache cache;
  if (blkid_get_cache(&cache, NULL) < 0) {
    LOG(INFO) << "blkid_get_cache failed";
    return fs_type;
  }
  blkid_dev dev = blkid_get_dev(cache, path.c_str(), BLKID_DEV_NORMAL);
  if (!dev) {
    LOG(INFO) << "blkid_get_dev failed";
    blkid_put_cache(cache);
    return fs_type;
  }

  const char *type, *value;
  blkid_tag_iterate iter = blkid_tag_iterate_begin(dev);
  while (blkid_tag_next(iter, &type, &value) == 0) {
    if (!strcmp(type, "TYPE")) {
      fs_type = value;
    }
  }
  blkid_tag_iterate_end(iter);
  blkid_put_cache(cache);
  return fs_type;
}

enum class DataImageAction { kNoAction, kCreateImage, kResizeImage };

static Result<DataImageAction> ChooseDataImageAction(
    const CuttlefishConfig::InstanceSpecific& instance) {
  if (instance.data_policy() == kDataPolicyAlwaysCreate) {
    return DataImageAction::kCreateImage;
  }
  if (!FileHasContent(instance.data_image())) {
    if (instance.data_policy() == kDataPolicyUseExisting) {
      return CF_ERR("A data image must exist to use -data_policy="
                    << kDataPolicyUseExisting);
    } else if (instance.data_policy() == kDataPolicyResizeUpTo) {
      return CF_ERR(instance.data_image()
                    << " does not exist, but resizing was requested");
    }
    return DataImageAction::kCreateImage;
  }
  if (instance.data_policy() == kDataPolicyUseExisting) {
    return DataImageAction::kNoAction;
  }
  auto current_fs_type = GetFsType(instance.data_image());
  if (current_fs_type != instance.userdata_format()) {
    CF_EXPECT(instance.data_policy() != kDataPolicyResizeUpTo,
              "Changing the fs format is incompatible with -data_policy="
                  << kDataPolicyResizeUpTo << " (\"" << current_fs_type
                  << "\" != \"" << instance.userdata_format() << "\")");
    return DataImageAction::kCreateImage;
  }
  if (instance.data_policy() == kDataPolicyResizeUpTo) {
    return DataImageAction::kResizeImage;
  }
  return DataImageAction::kNoAction;
}

} // namespace

Result<void> CreateBlankImage(const std::string& image, int num_mb,
                              const std::string& image_fmt) {
  LOG(DEBUG) << "Creating " << image;

  off_t image_size_bytes = static_cast<off_t>(num_mb) << 20;
  // The newfs_msdos tool with the mandatory -C option will do the same
  // as below to zero the image file, so we don't need to do it here
  if (image_fmt != "sdcard") {
    auto fd = SharedFD::Open(image, O_CREAT | O_TRUNC | O_RDWR, 0666);
    CF_EXPECTF(fd->Truncate(image_size_bytes) == 0,
               "`truncate --size={}M '{}'` failed: {}", num_mb, image,
               fd->StrError());
  }

  if (image_fmt == "ext4") {
    CF_EXPECT(Execute({"/sbin/mkfs.ext4", image}) == 0);
  } else if (image_fmt == "f2fs") {
    auto make_f2fs_path = HostBinaryPath("make_f2fs");
    CF_EXPECT(
        Execute({make_f2fs_path, "-l", "data", image, "-C", "utf8", "-O",
                 "compression,extra_attr,project_quota,casefold", "-g",
                 "android", "-b", F2FS_BLOCKSIZE, "-w", F2FS_BLOCKSIZE}) == 0);
  } else if (image_fmt == "sdcard") {
    // Reserve 1MB in the image for the MBR and padding, to simulate what
    // other OSes do by default when partitioning a drive
    off_t offset_size_bytes = 1 << 20;
    image_size_bytes -= offset_size_bytes;
    CF_EXPECT(NewfsMsdos(image, num_mb, 1), "Failed to create SD-Card fs");
    // Write the MBR after the filesystem is formatted, as the formatting tools
    // don't consistently preserve the image contents
    MasterBootRecord mbr = {
        .partitions = {{
            .partition_type = 0xC,
            .first_lba = (std::uint32_t)offset_size_bytes / kSectorSize,
            .num_sectors = (std::uint32_t)image_size_bytes / kSectorSize,
        }},
        .boot_signature = {0x55, 0xAA},
    };
    auto fd = SharedFD::Open(image, O_RDWR);
    CF_EXPECTF(WriteAllBinary(fd, &mbr) == sizeof(MasterBootRecord),
               "Writing MBR to '{}' failed: '{}'", image, fd->StrError());
  } else if (image_fmt != "none") {
    LOG(WARNING) << "Unknown image format '" << image_fmt
                 << "' for " << image << ", treating as 'none'.";
  }
  return {};
}

Result<void> InitializeDataImage(
    const CuttlefishConfig::InstanceSpecific& instance) {
  auto action = CF_EXPECT(ChooseDataImageAction(instance));
  switch (action) {
    case DataImageAction::kNoAction:
      LOG(DEBUG) << instance.data_image() << " exists. Not creating it.";
      return {};
    case DataImageAction::kCreateImage: {
      RemoveFile(instance.new_data_image());
      CF_EXPECT(instance.blank_data_image_mb() != 0,
                "Expected `-blank_data_image_mb` to be set for "
                "image creation.");
      CF_EXPECT(CreateBlankImage(instance.new_data_image(),
                                 instance.blank_data_image_mb(),
                                 instance.userdata_format()),
                "Failed to create a blank image at \""
                    << instance.new_data_image() << "\" with size "
                    << instance.blank_data_image_mb() << " and format \""
                    << instance.userdata_format() << "\"");
      return {};
    }
    case DataImageAction::kResizeImage: {
      CF_EXPECT(instance.blank_data_image_mb() != 0,
                "Expected `-blank_data_image_mb` to be set for "
                "image resizing.");
      CF_EXPECTF(Copy(instance.data_image(), instance.new_data_image()),
                 "Failed to `cp {} {}`", instance.data_image(),
                 instance.new_data_image());
      CF_EXPECT(ResizeImage(instance.new_data_image(),
                            instance.blank_data_image_mb(), instance),
                "Failed to resize \"" << instance.new_data_image() << "\" to "
                                      << instance.blank_data_image_mb()
                                      << " MB");
      return {};
    }
  }
}

Result<void> InitializeMiscImage(
    const CuttlefishConfig::InstanceSpecific& instance) {
  if (FileHasContent(instance.misc_image())) {
    LOG(DEBUG) << "misc partition image already exists";
    return {};
  }

  LOG(DEBUG) << "misc partition image: creating empty at \""
             << instance.misc_image() << "\"";
  CF_EXPECT(CreateBlankImage(instance.misc_image(), 1 /* mb */, "none"),
            "Failed to create misc image");
  return {};
}

static bool EspRequiredForBootFlow(BootFlow flow) {
  return flow == BootFlow::AndroidEfiLoader || flow == BootFlow::ChromeOs ||
         flow == BootFlow::Linux || flow == BootFlow::Fuchsia;
}

static bool EspRequiredForAPBootFlow(APBootFlow ap_boot_flow) {
  return ap_boot_flow == APBootFlow::Grub;
}

static void InitLinuxArgs(Arch target_arch, LinuxEspBuilder& linux) {
  linux.Root("/dev/vda2");

  linux.Argument("console", "hvc0").Argument("panic", "-1").Argument("noefi");

  switch (target_arch) {
    case Arch::Arm:
    case Arch::Arm64:
      linux.Argument("console", "ttyAMA0");
      break;
    case Arch::RiscV64:
      linux.Argument("console", "ttyS0");
      break;
    case Arch::X86:
    case Arch::X86_64:
      linux.Argument("console", "ttyS0")
          .Argument("pnpacpi", "off")
          .Argument("acpi", "noirq")
          .Argument("reboot", "k")
          .Argument("noexec", "off");
      break;
  }
}

static void InitChromeOsArgs(LinuxEspBuilder& linux) {
  linux.Root("/dev/vda2")
      .Argument("console", "ttyS0")
      .Argument("panic", "-1")
      .Argument("noefi")
      .Argument("init=/sbin/init")
      .Argument("boot=local")
      .Argument("rootwait")
      .Argument("noresume")
      .Argument("noswap")
      .Argument("loglevel=7")
      .Argument("noinitrd")
      .Argument("cros_efi")
      .Argument("cros_debug")
      .Argument("earlyprintk=serial,ttyS0,115200")
      .Argument("earlycon=uart8250,io,0x3f8")
      .Argument("pnpacpi", "off")
      .Argument("acpi", "noirq")
      .Argument("reboot", "k")
      .Argument("noexec", "off");
}

static bool BuildAPImage(const CuttlefishConfig& config,
                         const CuttlefishConfig::InstanceSpecific& instance) {
  auto linux = LinuxEspBuilder(instance.ap_esp_image_path());
  InitLinuxArgs(instance.target_arch(), linux);

  auto openwrt_args = OpenwrtArgsFromConfig(instance);
  for (auto& openwrt_arg : openwrt_args) {
    linux.Argument(openwrt_arg.first, openwrt_arg.second);
  }

  linux.Root("/dev/vda2")
      .Architecture(instance.target_arch())
      .Kernel(config.ap_kernel_image());

  return linux.Build();
}

static bool BuildOSImage(const CuttlefishConfig::InstanceSpecific& instance) {
  switch (instance.boot_flow()) {
    case BootFlow::AndroidEfiLoader: {
      auto android_efi_loader =
          AndroidEfiLoaderEspBuilder(instance.esp_image_path());
      android_efi_loader.EfiLoaderPath(instance.android_efi_loader())
          .Architecture(instance.target_arch());
      return android_efi_loader.Build();
    }
    case BootFlow::ChromeOs: {
      auto linux = LinuxEspBuilder(instance.esp_image_path());
      InitChromeOsArgs(linux);

      linux.Root("/dev/vda3")
          .Architecture(instance.target_arch())
          .Kernel(instance.chromeos_kernel_path());

      return linux.Build();
    }
    case BootFlow::Linux: {
      auto linux = LinuxEspBuilder(instance.esp_image_path());
      InitLinuxArgs(instance.target_arch(), linux);

      linux.Root("/dev/vda2")
          .Architecture(instance.target_arch())
          .Kernel(instance.linux_kernel_path());

      if (!instance.linux_initramfs_path().empty()) {
        linux.Initrd(instance.linux_initramfs_path());
      }

      return linux.Build();
    }
    case BootFlow::Fuchsia: {
      auto fuchsia = FuchsiaEspBuilder(instance.esp_image_path());
      return fuchsia.Architecture(instance.target_arch())
          .Zedboot(instance.fuchsia_zedboot_path())
          .MultibootBinary(instance.fuchsia_multiboot_bin_path())
          .Build();
    }
    default:
      break;
  }

  return true;
}

Result<void> InitializeEspImage(
    const CuttlefishConfig& config,
    const CuttlefishConfig::InstanceSpecific& instance) {
  if (EspRequiredForAPBootFlow(instance.ap_boot_flow())) {
    LOG(DEBUG) << "creating esp_image: " << instance.ap_esp_image_path();
    CF_EXPECT(BuildAPImage(config, instance));
  }
  const auto is_not_gem5 = config.vm_manager() != VmmMode::kGem5;
  const auto esp_required_for_boot_flow =
      EspRequiredForBootFlow(instance.boot_flow());
  if (is_not_gem5 && esp_required_for_boot_flow) {
    LOG(DEBUG) << "creating esp_image: " << instance.esp_image_path();
    CF_EXPECT(BuildOSImage(instance));
  }
  return {};
}

} // namespace cuttlefish