#!/usr/bin/env python

# Copyright 2019, 2020, 2022 VMware, Inc.
# All rights reserved.  -- VMware Confidential

"""Library to generate an ESX DD image.

ESX DD contains 3 VFAT partitions, each of which is sized to fit its content to
minimize the size of the overall .DD image:

  -  BOOT       contains the BIOS/UEFI bootloader files
  -  BOOTBANK1  contains the ESX modules
  -  BOOTBANK2  contains the locker files

When initializing a DD-imaged system, ESX is responsible for the following
tasks:

  1. Extend BOOT, BOOTBANK1, BOOTBANK2 to their actual runtime size
  2. Create the OSDATA partition (which is not present in the .DD)
  3. Move the locker files from BOOTBANK2 to their runtime location in OSDATA
  4. Initialize BOOTBANK2 as an empty bootbank.
"""
from math import ceil
import os
from tempfile import TemporaryDirectory

from systemStorage import BOOTBANK_SIZE_SMALL_MB
from systemStorage.autoPartition import addVfatPartitions
from systemStorage.esxboot import installBootloader
from systemStorage.esxdisk import EsxDD
from systemStorage.vfat import mcopy


def _iterFilesInDir(dirPath):
   """Iterate over all file paths under the given directory, recursively.
      Side effect: delete the dead symlinks.
   """
   for root, _, files in os.walk(dirPath):
      for basename in files:
         path = os.path.join(root, basename)
         if os.path.isfile(path):
            yield path
         elif os.path.islink(path) and not os.path.exists(path):
            os.unlink(path)

def _dirHasFiles(dirPath):
   """Returns True if dirPath contains at least one file, False otherwise.
   """
   for _ in _iterFilesInDir(dirPath):
      return True

   return False


class DDImage(EsxDD):
   """Represents an ESX DD image compatible with the 7.0+ parition layout.
   """

   def __init__(self, filePath, mmdBin=None, mcopyBin=None):
      self._mmdBin = mmdBin
      self._mcopyBin = mcopyBin
      super().__init__(filePath)

   @staticmethod
   def _getVfatPartSize(dirPath):
      """Compute the optimal size for the VFAT partitions in MB.

      XXX: Adding 2MB to accomodate VFAT internal metadata and fragmentation.
           This is empirical, we should re-visit this.
      """
      assert os.path.isdir(dirPath)
      BUFFER_IN_MB = 2
      sizeInBytes = sum([os.path.getsize(f) for f in _iterFilesInDir(dirPath)])
      return ceil(sizeInBytes / (1024 * 1024)) + BUFFER_IN_MB

   def _copyDirContentToPartition(self, directory, partition):
      """Copy the content of "directory" to the given partition but not
      "directory" itself.
      """
      if not _dirHasFiles(directory):
         return

      offset = partition.start * self.sectorSize
      files = [os.path.join(directory, basename) \
               for basename in os.listdir(directory)]
      mcopy(self.path, files, byteOffset=offset, exe=self._mcopyBin)

   def write(self, bootloaderDir, bootbank1Dir, lockerDir, hostSeedDir=None):
      """Write the DD image.
      """
      bootSize = self._getVfatPartSize(bootloaderDir)
      bb1Size = self._getVfatPartSize(bootbank1Dir)
      bb2Size = self._getVfatPartSize(lockerDir)
      if hostSeedDir:
         bb2Size = bb2Size + self._getVfatPartSize(hostSeedDir)

      sectorsPerMB = 1024 * 1024 // self.sectorSize

      # create or truncate the .dd output file
      open(self.path, 'w').close()

      # create partitions
      parts = addVfatPartitions(self, BOOTBANK_SIZE_SMALL_MB, sectorsPerMB,
                                bootPartSize=bootSize, bootbank1Size=bb1Size,
                                bootbank2Size=bb2Size)

      self.syncPartitions(autoFormat=True, skipBackupGpt=True)

      installBootloader(self, srcRoot=bootloaderDir, mcopyBin=self._mcopyBin,
                        mmdBin=self._mmdBin)

      self._copyDirContentToPartition(bootbank1Dir, parts[1])
      self._copyDirContentToPartition(lockerDir, parts[2])
      if hostSeedDir:
         self._copyDirContentToPartition(hostSeedDir, parts[2])
