// Copyright 2017 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/param.h>
#include <zircon/assert.h>
#include <zircon/device/block.h>
#include <zircon/errors.h>
#include <zircon/syscalls.h>
#include <zircon/types.h>

#include <chromeos-disk-setup/chromeos-disk-setup.h>
#include <gpt/cros.h>
#include <gpt/gpt.h>

namespace chromeos_disk_setup {
using gpt::GptDevice;
namespace {

constexpr uint8_t kFvmGuid[GPT_GUID_LEN] = GUID_FVM_VALUE;
constexpr uint8_t kKernGuid[GPT_GUID_LEN] = GUID_CROS_KERNEL_VALUE;
constexpr uint8_t kRootGuid[GPT_GUID_LEN] = GUID_CROS_ROOT_VALUE;
constexpr uint8_t kStateGuid[GPT_GUID_LEN] = GUID_CROS_STATE_VALUE;
constexpr uint8_t kSysCfgGuid[GPT_GUID_LEN] = GUID_SYS_CONFIG_VALUE;

// this value is shared with device-partitioner.cpp
constexpr uint64_t kMinFvmSize = 8LU * (1 << 30);

constexpr uint64_t kSysCfgSize = 1 << 20;

// part_name_eql compares the name field in the given partition to the cstring
// name, returning true if the partition name is equal to name and zero-padded.
bool part_name_eql(const gpt_partition_t* part, const char* name) {
  if (part == NULL) {
    return false;
  }
  char buf[GPT_NAME_LEN] = {0};
  utf16_to_cstring(&buf[0], (const uint16_t*)part->name, GPT_NAME_LEN / 2);
  return !strncmp(buf, name, GPT_NAME_LEN);
}

// part_name_guid_eql returns true if the given partition has the given name and
// the given type GUID, false otherwise.
bool part_name_guid_eql(const gpt_partition_t* part, const char* name,
                        const uint8_t guid[GPT_GUID_LEN]) {
  if (part == NULL) {
    return false;
  }
  if (!memcmp(part->type, &guid[0], GPT_GUID_LEN) && part_name_eql(part, name)) {
    return true;
  }
  return false;
}

// part_size_gte returns true if the partition size is greater than or equal to
// the size given, false otherwise.
bool part_size_gte(gpt_partition_t* part, uint64_t size, uint64_t block_size) {
  if (part == NULL) {
    return false;
  }
  uint64_t size_in_blocks = part->last - part->first + 1;
  return size_in_blocks * block_size >= size;
}

// find_by_type finds the first partition matching the given type guid.
gpt_partition_t* find_by_type(const GptDevice* gpt, const uint8_t type_guid[GPT_GUID_LEN]) {
  for (uint32_t i = 0; i < gpt::kPartitionCount; ++i) {
    gpt_partition_t* p = gpt->GetPartition(i);
    if (p == NULL) {
      continue;
    }
    if (!memcmp(p->type, &type_guid[0], GPT_GUID_LEN)) {
      return p;
    }
  }
  return NULL;
}

// find_by_type_and_name finds the first partition matching the given type guid and name.
gpt_partition_t* find_by_type_and_name(const GptDevice* gpt, const uint8_t type_guid[GPT_GUID_LEN],
                                       const char* name) {
  for (uint32_t i = 0; i < gpt::kPartitionCount; ++i) {
    gpt_partition_t* p = gpt->GetPartition(i);
    if (p == NULL) {
      continue;
    }
    if (part_name_guid_eql(p, name, type_guid)) {
      return p;
    }
  }
  return NULL;
}

// Find a contiguous run of free space on the disk at least blocks_req in length.
// If space is found, true is returned and out_hole_start and out_hole_end
// contain the first free and last free blocks in a contiguous run.
bool find_space(GptDevice* gpt, const uint64_t blocks_req, uint64_t* out_hole_start,
                uint64_t* out_hole_end) {
  gpt_partition_t* parts[gpt::kPartitionCount] = {0};
  for (uint32_t i = 0; i < gpt::kPartitionCount; i++) {
    parts[i] = gpt->GetPartition(i);
  }

  // XXX(raggi): once the lib supports more and less than gpt::kPartitionCount,
  // this will need to be fixed (as will many loops).
  gpt_sort_partitions(parts, gpt::kPartitionCount);

  uint64_t first_usable, last_usable;
  ZX_ASSERT(gpt->Range(&first_usable, &last_usable) == ZX_OK);

  uint64_t prev_end = first_usable - 1;
  for (uint32_t i = 0; i < gpt::kPartitionCount; i++) {
    if (parts[i] == NULL) {
      continue;
    }

    // TODO(raggi): find out how the tests end up making this state
    if (parts[i]->first >= last_usable || parts[i]->last >= last_usable) {
      break;
    }

    uint64_t gap = parts[i]->first - (prev_end + 1);
    if (gap >= blocks_req) {
      *out_hole_start = prev_end + 1;
      *out_hole_end = parts[i]->first - 1;

      return true;
    }
    prev_end = parts[i]->last;
  }

  if (prev_end < last_usable && last_usable - (prev_end + 1) >= blocks_req) {
    *out_hole_start = prev_end + 1;
    *out_hole_end = last_usable;

    return true;
  }

  return false;
}

// create a GPT entry with the supplied attributes and assign it a random
// GUID. May return an error if the GUID can not be generated for the partition
// or operations on the gpt_device_t fail.
zx_status_t create_gpt_entry(GptDevice* gpt, const uint64_t first, const uint64_t blks,
                             const uint8_t type[GPT_GUID_LEN], const char* name) {
  uint8_t guid[GPT_GUID_LEN];
  zx_cprng_draw(guid, GPT_GUID_LEN);

  uint8_t tguid[GPT_GUID_LEN];
  memcpy(&tguid, type, GPT_GUID_LEN);
  if (gpt->AddPartition(name, tguid, guid, first, blks, 0) != ZX_OK) {
    return ZX_ERR_INTERNAL;
  }

  return ZX_OK;
}

}  // namespace

__BEGIN_CDECLS

bool is_cros(const GptDevice* gpt) {
  uint8_t roots = 0;
  uint8_t kerns = 0;
  bool state = false;

  for (uint32_t i = 0; i < gpt::kPartitionCount; ++i) {
    gpt_partition_t* p = gpt->GetPartition(i);
    if (p == NULL) {
      continue;
    }

    if (!memcmp(p->type, kRootGuid, GPT_GUID_LEN) &&
        (part_name_eql(p, "ROOT-A") || part_name_eql(p, "ROOT-B"))) {
      roots++;
    } else if (!memcmp(p->type, kKernGuid, GPT_GUID_LEN) &&
               (part_name_eql(p, "KERN-A") || part_name_eql(p, "KERN-B"))) {
      kerns++;
    } else if (!memcmp(p->type, kStateGuid, GPT_GUID_LEN) && part_name_eql(p, "STATE")) {
      state = true;
    }
  }

  return state && roots >= 2 && kerns >= 2;
}

// is_ready_to_pave returns true if there exist partitions for:
// * ZIRCON-A is a GUID_CROS_KERNEL_VALUE at least sz_kern in size.
// * ZIRCON-B is a GUID_CROS_KERNEL_VALUE at least sz_kern in size.
// * ZIRCON-R is a GUID_CROS_KERNEL_VALUE at least sz_kern in size.
// * FVM      is a GUID_FVM_VALUE         at least sz_kern * 8 in size
bool is_ready_to_pave(const GptDevice* gpt, const fuchsia_hardware_block_BlockInfo* blk_info,
                      const uint64_t sz_kern) {
  bool found_zircon_a = false, found_zircon_b = false, found_zircon_r = false;
  bool found_fvm = false, found_syscfg = false;

  for (uint32_t i = 0; i < gpt::kPartitionCount; i++) {
    gpt_partition_t* part = gpt->GetPartition(i);
    if (part == NULL) {
      continue;
    }
    if (!memcmp(part->type, kFvmGuid, GPT_GUID_LEN)) {
      if (!part_size_gte(part, kMinFvmSize, blk_info->block_size)) {
        continue;
      }
      found_fvm = true;
      continue;
    }
    if (!memcmp(part->type, kKernGuid, GPT_GUID_LEN)) {
      if (!part_size_gte(part, sz_kern, blk_info->block_size)) {
        continue;
      }
      if (part_name_eql(part, "ZIRCON-A")) {
        found_zircon_a = true;
      }
      if (part_name_eql(part, "ZIRCON-B")) {
        found_zircon_b = true;
      }
      if (part_name_eql(part, "ZIRCON-R")) {
        found_zircon_r = true;
      }
    }
    if (!memcmp(part->type, kSysCfgGuid, GPT_GUID_LEN)) {
      if (!part_size_gte(part, kSysCfgSize, blk_info->block_size)) {
        continue;
      }
      found_syscfg = true;
    }
  }
  if (!found_syscfg) {
    printf("cros-disk-setup: missing syscfg (or insufficient size)\n");
  }
  if (!found_fvm) {
    printf("cros-disk-setup: missing FVM (or insufficient size)\n");
  }
  if (!found_zircon_a || !found_zircon_b || !found_zircon_r) {
    printf("cros-disk-setup: missing one or more kernel partitions\n");
  }

  return found_zircon_a && found_zircon_b && found_zircon_r && found_fvm && found_syscfg;
}

zx_status_t config_cros_for_fuchsia(GptDevice* gpt,
                                    const fuchsia_hardware_block_BlockInfo* blk_info,
                                    const uint64_t sz_kern) {
  // TODO(raggi): this ends up getting called twice, as the canonical user,
  // the paver, calls is_ready_to_pave itself in order to determine first
  // whether it will need to sync the gpt.
  if (is_ready_to_pave(gpt, blk_info, sz_kern)) {
    return ZX_OK;
  }

  // TODO(ZX-1396): The gpt_device_t may not be valid for modification if it
  // is a newly initialized GPT which has never had gpt_device_finalize or
  // gpt_device_sync called.
  if (gpt->Finalize() != ZX_OK) {
    return ZX_ERR_INTERNAL;
  }

  // Remove the pre-existing Fuchsia partitions as when we were not already
  // pave-able and we're paving, assume that we want to tend toward a golden
  // layout. This also avoids any additional complexity that could arise from
  // intermediate gaps between these partitions.

  gpt_partition_t* p;
  if ((p = find_by_type_and_name(gpt, kKernGuid, "ZIRCON-A")) != NULL) {
    ZX_ASSERT(gpt->RemovePartition(p->guid) == ZX_OK);
  }
  if ((p = find_by_type_and_name(gpt, kKernGuid, "ZIRCON-B")) != NULL) {
    ZX_ASSERT(gpt->RemovePartition(p->guid) == ZX_OK);
  }
  if ((p = find_by_type_and_name(gpt, kKernGuid, "ZIRCON-R")) != NULL) {
    ZX_ASSERT(gpt->RemovePartition(p->guid) == ZX_OK);
  }
  if ((p = find_by_type(gpt, kFvmGuid)) != NULL) {
    ZX_ASSERT(gpt->RemovePartition(p->guid) == ZX_OK);
  }
  if ((p = find_by_type_and_name(gpt, kSysCfgGuid, "SYSCFG")) != NULL) {
    ZX_ASSERT(gpt->RemovePartition(p->guid) == ZX_OK);
  }

  // Space is required for 3 kernel partitions and one FVM partition that is
  // at least 8 kernels in size.
  uint64_t needed_space = sz_kern * 3 + kMinFvmSize + kSysCfgSize;

  // see if a contiguous block of space is available for space needed
  uint64_t blocks_needed = howmany(needed_space, blk_info->block_size);

  uint64_t hole_start, hole_end;
  bool found_hole = find_space(gpt, blocks_needed, &hole_start, &hole_end);

  // TODO(raggi): find a good heuristic to detect "old-paver" behavior, and if
  // we can detect that, remove the -C's, otherwise leave them alone.

  // First try removing the kernc and rootc partitions, as they're often a good fit for us:
  if (!found_hole) {
    // Some partitions were not large enough. If we found a KERN-C or a ROOT-C, delete them:
    if ((p = find_by_type_and_name(gpt, kKernGuid, "KERN-C")) != NULL) {
      ZX_ASSERT(gpt->RemovePartition(p->guid) == ZX_OK);
    }
    if ((p = find_by_type_and_name(gpt, kRootGuid, "ROOT-C")) != NULL) {
      ZX_ASSERT(gpt->RemovePartition(p->guid) == ZX_OK);
    }

    found_hole = find_space(gpt, blocks_needed, &hole_start, &hole_end);
  }

  // Still not enough contiguous space is available on disk, try shrinking STATE
  if (!found_hole && (p = find_by_type_and_name(gpt, kStateGuid, "STATE")) != NULL) {
    uint64_t min_state_sz_blks = howmany(MIN_SZ_STATE, blk_info->block_size);

    // TODO (TO-607) consider if there is free space on either side of STATE

    // The STATE partition is expected to be at the end of the GPT in cros,
    // and can be shrunk in order to make space for use cases such as this.
    // Here we first try to make the partition half of it's current size
    // (pinned to a minimum of min_sz_blks). This gives us a roughly equal
    // share of the free space on the disk.

    uint64_t state_part_blks = (p->last - p->first) + 1;
    uint64_t new_state_part_blks = state_part_blks / 2;
    if (new_state_part_blks < min_state_sz_blks) {
      new_state_part_blks = min_state_sz_blks;
    }
    uint64_t new_free_blks = state_part_blks - new_state_part_blks;

    if (new_free_blks >= blocks_needed) {
      p->first = p->first + new_free_blks;

      // find_space is re-run here as there is often a chunk of free space
      // left before the STATE partition that is not big enough for us to
      // fit, but is sensible to use which would not be used simply by
      // using the state->first offset.
      found_hole = find_space(gpt, blocks_needed, &hole_start, &hole_end);
    }
  }

  if (!found_hole) {
    return ZX_ERR_NO_SPACE;
  }

  printf("cros-disk-setup: Creating SYSCFG\n");
  zx_status_t status;
  uint64_t sz_syscfg_blks = howmany(kSysCfgSize, blk_info->block_size);
  if ((status = create_gpt_entry(gpt, hole_start, sz_syscfg_blks, kSysCfgGuid, "SYSCFG")) !=
      ZX_OK) {
    printf("cros-disk-setup: Error creating SYSCFG: %d\n", status);
    return ZX_ERR_INTERNAL;
  }
  hole_start += sz_syscfg_blks;

  uint64_t sz_kern_blks = howmany(sz_kern, blk_info->block_size);

  // Create GPT entries for ZIRCON-A, ZIRCON-B, ZIRCON-R and FVM if needed.
  const char* kernel_names[3] = {"ZIRCON-A", "ZIRCON-B", "ZIRCON-R"};
  for (size_t i = 0; i < 3; ++i) {
    printf("cros-disk-setup: Creating %s\n", kernel_names[i]);
    if ((status = create_gpt_entry(gpt, hole_start, sz_kern_blks, kKernGuid, kernel_names[i])) !=
        ZX_OK) {
      printf("cros-disk-setup: Error creating %s: %d\n", kernel_names[i], status);
      return ZX_ERR_INTERNAL;
    }
    hole_start += sz_kern_blks;
  }

  printf("cros-disk-setup: Creating FVM\n");

  // TODO(raggi): add this after the test setup supports it.
  // // clear the FVM superblock to ensure that a new FVM will be created there.
  // if (gpt_partition_clear(gpt, hole_start, 1) != 0) {
  //     printf("Error clearing FVM superblock.\n");
  //     return ZX_ERR_INTERNAL;
  // }

  // The created FVM partition will fill the available free space.
  if ((status = create_gpt_entry(gpt, hole_start, (hole_end - hole_start), kFvmGuid, "fvm")) !=
      ZX_OK) {
    printf("cros-disk-setup: Error creating FVM\n");
    return status;
  }

  // TODO(raggi): add this once the test setup supports it.
  // if (!gpt_device_finalize(gpt)) {
  //     printf("Error finalizing GPT\n");
  //     return ZX_ERR_INTERNAL;
  // }
  return ZX_OK;
}

__END_CDECLS

}  // namespace chromeos_disk_setup